fix: copy to Clipboard order (#16299)

* copy to Clipboard order

* centralized copyToClipboard

* fixed table order

* fixed tests

* added colnames to all viz types

* added colnames to all viz types

* added colnames to all viz types

superset-frontend/src/SqlLab/components/ResultSet.tsx

@@ -448,7 +448,7 @@ export default class ResultSet extends React.PureComponent<
       if (this.props.cache && this.props.query.cached) {
         ({ data } = this.state);
       }
-
+      const { columns } = this.props.query.results;
       // Added compute logic to stop user from being able to Save & Explore
       const {
         saveDatasetRadioBtnState,
@@ -508,7 +508,7 @@ export default class ResultSet extends React.PureComponent<
             )}
 
             <CopyToClipboard
-              text={prepareCopyToClipboardTabularData(data)}
+              text={prepareCopyToClipboardTabularData(data, columns)}
               wrapped={false}
               copyNode={
                 <Button buttonSize="small">

superset-frontend/src/components/TableView/TableView.tsx

@@ -158,7 +158,6 @@ const TableView = ({
     useSortBy,
     usePagination,
   );
-
   useEffect(() => {
     if (serverPagination && pageIndex !== initialState.pageIndex) {
       onServerPagination({

superset-frontend/src/explore/components/DataTableControl/index.tsx

@@ -55,11 +55,15 @@ const CopyNode = (
 
 export const CopyToClipboardButton = ({
   data,
+  columns,
 }: {
   data?: Record<string, any>;
+  columns?: string[];
 }) => (
   <CopyToClipboard
-    text={data ? prepareCopyToClipboardTabularData(data) : ''}
+    text={
+      data && columns ? prepareCopyToClipboardTabularData(data, columns) : ''
+    }
     wrapped={false}
     copyNode={CopyNode}
   />
@@ -113,29 +117,32 @@ export const useFilteredTableData = (
   }, [data, filterText]);
 
 export const useTableColumns = (
+  colnames?: string[],
   data?: Record<string, any>[],
   moreConfigs?: { [key: string]: Partial<Column> },
 ) =>
   useMemo(
     () =>
-      data?.length
-        ? Object.keys(data[0]).map(
-            key =>
-              ({
-                accessor: row => row[key],
-                Header: key,
-                Cell: ({ value }) => {
-                  if (value === true) {
-                    return BOOL_TRUE_DISPLAY;
-                  }
-                  if (value === false) {
-                    return BOOL_FALSE_DISPLAY;
-                  }
-                  return String(value);
-                },
-                ...moreConfigs?.[key],
-              } as Column),
-          )
+      colnames && data?.length
+        ? colnames
+            .filter((column: string) => Object.keys(data[0]).includes(column))
+            .map(
+              key =>
+                ({
+                  accessor: row => row[key],
+                  Header: key,
+                  Cell: ({ value }) => {
+                    if (value === true) {
+                      return BOOL_TRUE_DISPLAY;
+                    }
+                    if (value === false) {
+                      return BOOL_FALSE_DISPLAY;
+                    }
+                    return String(value);
+                  },
+                  ...moreConfigs?.[key],
+                } as Column),
+            )
         : [],
     [data, moreConfigs],
   );

superset-frontend/src/explore/components/DataTableControl/useTableColumns.test.ts

@@ -42,9 +42,10 @@ const data = [
     [unicodeKey]: unicodeKey,
   },
 ];
+const all_columns = ['col01', 'col02', 'col03', asciiKey, unicodeKey];
 
 test('useTableColumns with no options', () => {
-  const hook = renderHook(() => useTableColumns(data));
+  const hook = renderHook(() => useTableColumns(all_columns, data));
   expect(hook.result.current).toEqual([
     {
       Cell: expect.any(Function),
@@ -83,7 +84,7 @@ test('useTableColumns with no options', () => {
 
 test('use only the first record columns', () => {
   const newData = [data[3], data[0]];
-  const hook = renderHook(() => useTableColumns(newData));
+  const hook = renderHook(() => useTableColumns(all_columns, newData));
   expect(hook.result.current).toEqual([
     {
       Cell: expect.any(Function),
@@ -134,7 +135,9 @@ test('use only the first record columns', () => {
 });
 
 test('useTableColumns with options', () => {
-  const hook = renderHook(() => useTableColumns(data, { col01: { id: 'ID' } }));
+  const hook = renderHook(() =>
+    useTableColumns(all_columns, data, { col01: { id: 'ID' } }),
+  );
   expect(hook.result.current).toEqual([
     {
       Cell: expect.any(Function),

superset-frontend/src/explore/components/DataTablesPane/DataTablesPane.test.tsx

@@ -58,6 +58,11 @@ const createProps = () => ({
   tableSectionHeight: 156.9,
   chartStatus: 'rendered',
   onCollapseChange: jest.fn(),
+  queriesResponse: [
+    {
+      colnames: [],
+    },
+  ],
 });
 
 afterAll(() => {

superset-frontend/src/explore/components/DataTablesPane/index.tsx

@@ -112,6 +112,7 @@ export const DataTablesPane = ({
   chartStatus,
   ownState,
   errorMessage,
+  queriesResponse,
 }: {
   queryFormData: Record<string, any>;
   tableSectionHeight: number;
@@ -119,6 +120,7 @@ export const DataTablesPane = ({
   ownState?: JsonObject;
   onCollapseChange: (openPanelName: string) => void;
   errorMessage?: JSX.Element;
+  queriesResponse: Record<string, any>;
 }) => {
   const [data, setData] = useState<{
     [RESULT_TYPES.results]?: Record<string, any>[];
@@ -128,6 +130,7 @@ export const DataTablesPane = ({
     [RESULT_TYPES.results]: true,
     [RESULT_TYPES.samples]: true,
   });
+  const [columnNames, setColumnNames] = useState<string[]>([]);
   const [error, setError] = useState(NULLISH_RESULTS_STATE);
   const [filterText, setFilterText] = useState('');
   const [activeTabKey, setActiveTabKey] = useState<string>(
@@ -220,6 +223,13 @@ export const DataTablesPane = ({
     }));
   }, [queryFormData.adhoc_filters, queryFormData.datasource]);
 
+  useEffect(() => {
+    if (queriesResponse) {
+      const { colnames } = queriesResponse[0];
+      setColumnNames([...colnames]);
+    }
+  }, [queriesResponse]);
+
   useEffect(() => {
     if (panelOpen && isRequestPending[RESULT_TYPES.results]) {
       if (errorMessage) {
@@ -277,9 +287,17 @@ export const DataTablesPane = ({
     ),
   };
 
+  // this is to preserve the order of the columns, even if there are integer values,
+  // while also only grabbing the first column's keys
   const columns = {
-    [RESULT_TYPES.results]: useTableColumns(data[RESULT_TYPES.results]),
-    [RESULT_TYPES.samples]: useTableColumns(data[RESULT_TYPES.samples]),
+    [RESULT_TYPES.results]: useTableColumns(
+      columnNames,
+      data[RESULT_TYPES.results],
+    ),
+    [RESULT_TYPES.samples]: useTableColumns(
+      columnNames,
+      data[RESULT_TYPES.samples],
+    ),
   };
 
   const renderDataTable = (type: string) => {
@@ -316,7 +334,7 @@ export const DataTablesPane = ({
   const TableControls = (
     <TableControlsWrapper>
       <RowCount data={data[activeTabKey]} loading={isLoading[activeTabKey]} />
-      <CopyToClipboardButton data={data[activeTabKey]} />
+      <CopyToClipboardButton data={data[activeTabKey]} columns={columnNames} />
       <FilterInput onChangeHandler={setFilterText} />
     </TableControlsWrapper>
   );

superset-frontend/src/explore/components/ExploreChartPanel.jsx

@@ -116,7 +116,6 @@ const ExploreChartPanel = props => {
   const theme = useTheme();
   const gutterMargin = theme.gridUnit * GUTTER_SIZE_FACTOR;
   const gutterHeight = theme.gridUnit * GUTTER_SIZE_FACTOR;
-
   const { height: hHeight, ref: headerRef } = useResizeDetector({
     refreshMode: 'debounce',
     refreshRate: 300,
@@ -128,7 +127,6 @@ const ExploreChartPanel = props => {
   const [splitSizes, setSplitSizes] = useState(
     getFromLocalStorage(STORAGE_KEYS.sizes, INITIAL_SIZES),
   );
-
   const { slice } = props;
   const updateQueryContext = useCallback(
     async function fetchChartData() {
@@ -211,7 +209,6 @@ const ExploreChartPanel = props => {
     }
     setSplitSizes(splitSizes);
   };
-
   const renderChart = useCallback(() => {
     const { chart, vizType } = props;
     const newHeight =
@@ -317,6 +314,7 @@ const ExploreChartPanel = props => {
             onCollapseChange={onCollapseChange}
             chartStatus={props.chart.chartStatus}
             errorMessage={props.errorMessage}
+            queriesResponse={props.chart.queriesResponse}
           />
         </Split>
       )}

superset-frontend/src/utils/common.js

@@ -87,10 +87,21 @@ export function optionFromValue(opt) {
   return { value: optionValue(opt), label: optionLabel(opt) };
 }
 
-export function prepareCopyToClipboardTabularData(data) {
+export function prepareCopyToClipboardTabularData(data, columns) {
   let result = '';
   for (let i = 0; i < data.length; i += 1) {
-    result += `${Object.values(data[i]).join('\t')}\n`;
+    const row = {};
+    for (let j = 0; j < columns.length; j += 1) {
+      // JavaScript does not mantain the order of a mixed set of keys (i.e integers and strings)
+      // the below function orders the keys based on the column names.
+      const key = columns[j].name || columns[j];
+      if (data[i][key]) {
+        row[j] = data[i][key];
+      } else {
+        row[j] = data[i][parseFloat(key)];
+      }
+    }
+    result += `${Object.values(row).join('\t')}\n`;
   }
   return result;
 }

superset-frontend/src/utils/common.test.jsx

@@ -46,15 +46,18 @@ describe('utils/common', () => {
   describe('prepareCopyToClipboardTabularData', () => {
     it('converts empty array', () => {
       const array = [];
-      expect(prepareCopyToClipboardTabularData(array)).toEqual('');
+      const column = [];
+      expect(prepareCopyToClipboardTabularData(array, column)).toEqual('');
     });
     it('converts non empty array', () => {
       const array = [
         { column1: 'lorem', column2: 'ipsum' },
         { column1: 'dolor', column2: 'sit', column3: 'amet' },
       ];
-      expect(prepareCopyToClipboardTabularData(array)).toEqual(
-        'lorem\tipsum\ndolor\tsit\tamet\n',
+      const column = ['column1', 'column2', 'column3'];
+      console.log(prepareCopyToClipboardTabularData(array, column));
+      expect(prepareCopyToClipboardTabularData(array, column)).toEqual(
+        'lorem\tipsum\t\ndolor\tsit\tamet\n',
       );
     });
   });

superset/viz.py

@@ -450,7 +450,8 @@ class BaseViz:
 
         payload = self.get_df_payload(query_obj)
 
-        df = payload.get("df")
+        # if payload does not have a df, we are raising an error here.
+        df = cast(Optional[pd.DataFrame], payload["df"])
 
         if self.status != utils.QueryStatus.FAILED:
             payload["data"] = self.get_data(df)
@@ -482,7 +483,8 @@ class BaseViz:
             for col in filter_columns
             if col not in columns and col not in filter_values_columns
         ] + rejected_time_columns
-
+        if df is not None:
+            payload["colnames"] = list(df.columns)
         return payload
 
     def get_df_payload(