From e6ff25c9802d5a8b523a966a66168b739a97b476 Mon Sep 17 00:00:00 2001
From: Phillip Kelley-Dotson <pkelleydotson@yahoo.com>
Date: Fri, 12 Nov 2021 14:23:46 -0800
Subject: [PATCH] fix(chart): ensure samples data is displayed (#16900)

* initial fix

* use data keys for cols

* add e2e test

* remove code

* trim logic
---
 .../integration/explore/control.test.ts       |  19 +++
 .../components/DataTableControl/index.tsx     |   2 +-
 .../components/DataTablesPane/index.tsx       | 152 ++++++++++--------
 3 files changed, 109 insertions(+), 64 deletions(-)

diff --git a/superset-frontend/cypress-base/cypress/integration/explore/control.test.ts b/superset-frontend/cypress-base/cypress/integration/explore/control.test.ts
index c57369e891..b316d5c12e 100644
--- a/superset-frontend/cypress-base/cypress/integration/explore/control.test.ts
+++ b/superset-frontend/cypress-base/cypress/integration/explore/control.test.ts
@@ -113,6 +113,25 @@ describe('VizType control', () => {
   });
 });
 
+describe('Test datatable', () => {
+  beforeEach(() => {
+    cy.login();
+    interceptChart({ legacy: false }).as('tableChartData');
+    interceptChart({ legacy: true }).as('lineChartData');
+  });
+  it('Data Pane opens and loads results', () => {
+    cy.get('[data-test="data-tab"]').click();
+    cy.get('[data-test="row-count-label"]').contains('27 rows retrieved');
+    cy.contains('View results');
+    cy.get('.ant-empty-description').should('not.exist');
+  });
+  it('Datapane loads view samples', () => {
+    cy.contains('View samples').click();
+    cy.get('[data-test="row-count-label"]').contains('1k rows retrieved');
+    cy.get('.ant-empty-description').should('not.exist');
+  });
+});
+
 describe('Time range filter', () => {
   beforeEach(() => {
     cy.login();
diff --git a/superset-frontend/src/explore/components/DataTableControl/index.tsx b/superset-frontend/src/explore/components/DataTableControl/index.tsx
index d36a313020..a1383cf4a9 100644
--- a/superset-frontend/src/explore/components/DataTableControl/index.tsx
+++ b/superset-frontend/src/explore/components/DataTableControl/index.tsx
@@ -150,5 +150,5 @@ export const useTableColumns = (
                 } as Column),
             )
         : [],
-    [data, moreConfigs],
+    [data, colnames, moreConfigs],
   );
diff --git a/superset-frontend/src/explore/components/DataTablesPane/index.tsx b/superset-frontend/src/explore/components/DataTablesPane/index.tsx
index 605c4edf7b..67a07f905f 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/index.tsx
+++ b/superset-frontend/src/explore/components/DataTablesPane/index.tsx
@@ -106,6 +106,58 @@ const Error = styled.pre`
   margin-top: ${({ theme }) => `${theme.gridUnit * 4}px`};
 `;
 
+interface DataTableProps {
+  columnNames: string[];
+  filterText: string;
+  data: object[] | undefined;
+  isLoading: boolean;
+  error: string | undefined;
+  errorMessage: React.ReactElement | undefined;
+}
+
+const DataTable = ({
+  columnNames,
+  filterText,
+  data,
+  isLoading,
+  error,
+  errorMessage,
+}: DataTableProps) => {
+  // 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 = useTableColumns(columnNames, data);
+  const filteredData = useFilteredTableData(filterText, data);
+
+  if (isLoading) {
+    return <Loading />;
+  }
+  if (error) {
+    return <Error>{error}</Error>;
+  }
+  if (data) {
+    if (data.length === 0) {
+      return <span>No data</span>;
+    }
+    return (
+      <TableView
+        columns={columns}
+        data={filteredData}
+        pageSize={DATA_TABLE_PAGE_SIZE}
+        noDataText={t('No data')}
+        emptyWrapperType={EmptyWrapperType.Small}
+        className="table-condensed"
+        isPaginationSticky
+        showRowCount={false}
+        small
+      />
+    );
+  }
+  if (errorMessage) {
+    return <Error>{errorMessage}</Error>;
+  }
+  return null;
+};
+
 export const DataTablesPane = ({
   queryFormData,
   tableSectionHeight,
@@ -131,7 +183,13 @@ export const DataTablesPane = ({
     [RESULT_TYPES.results]: true,
     [RESULT_TYPES.samples]: true,
   });
-  const [columnNames, setColumnNames] = useState<string[]>([]);
+  const [columnNames, setColumnNames] = useState<{
+    [RESULT_TYPES.results]: string[];
+    [RESULT_TYPES.samples]: string[];
+  }>({
+    [RESULT_TYPES.results]: [],
+    [RESULT_TYPES.samples]: [],
+  });
   const [error, setError] = useState(NULLISH_RESULTS_STATE);
   const [filterText, setFilterText] = useState('');
   const [activeTabKey, setActiveTabKey] = useState<string>(
@@ -192,7 +250,13 @@ export const DataTablesPane = ({
               [resultType]: json.result[0].data,
             }));
           }
-
+          const checkCols = json?.result[0]?.data?.length
+            ? Object.keys(json.result[0].data[0])
+            : null;
+          setColumnNames(prevColumnNames => ({
+            ...prevColumnNames,
+            [resultType]: json.result[0].columns || checkCols,
+          }));
           setIsLoading(prevIsLoading => ({
             ...prevIsLoading,
             [resultType]: false,
@@ -215,7 +279,7 @@ export const DataTablesPane = ({
           });
         });
     },
-    [queryFormData],
+    [queryFormData, columnNames],
   );
 
   useEffect(() => {
@@ -239,7 +303,10 @@ export const DataTablesPane = ({
   useEffect(() => {
     if (queriesResponse && chartStatus === 'success') {
       const { colnames } = queriesResponse[0];
-      setColumnNames([...colnames]);
+      setColumnNames({
+        ...columnNames,
+        [RESULT_TYPES.results]: [...colnames],
+      });
     }
   }, [queriesResponse]);
 
@@ -289,67 +356,12 @@ export const DataTablesPane = ({
     errorMessage,
   ]);
 
-  const filteredData = {
-    [RESULT_TYPES.results]: useFilteredTableData(
-      filterText,
-      formattedData[RESULT_TYPES.results],
-    ),
-    [RESULT_TYPES.samples]: useFilteredTableData(
-      filterText,
-      formattedData[RESULT_TYPES.samples],
-    ),
-  };
-
-  // 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(
-      columnNames,
-      data[RESULT_TYPES.results],
-    ),
-    [RESULT_TYPES.samples]: useTableColumns(
-      columnNames,
-      data[RESULT_TYPES.samples],
-    ),
-  };
-
-  const renderDataTable = (type: string) => {
-    if (isLoading[type]) {
-      return <Loading />;
-    }
-    if (error[type]) {
-      return <Error>{error[type]}</Error>;
-    }
-    if (data[type]) {
-      if (data[type]?.length === 0) {
-        return <span>No data</span>;
-      }
-      return (
-        <TableView
-          columns={columns[type]}
-          data={filteredData[type]}
-          pageSize={DATA_TABLE_PAGE_SIZE}
-          noDataText={t('No data')}
-          emptyWrapperType={EmptyWrapperType.Small}
-          className="table-condensed"
-          isPaginationSticky
-          showRowCount={false}
-          small
-        />
-      );
-    }
-    if (errorMessage) {
-      return <Error>{errorMessage}</Error>;
-    }
-    return null;
-  };
-
   const TableControls = (
     <TableControlsWrapper>
       <RowCount data={data[activeTabKey]} loading={isLoading[activeTabKey]} />
       <CopyToClipboardButton
         data={formattedData[activeTabKey]}
-        columns={columnNames}
+        columns={columnNames[activeTabKey]}
       />
       <FilterInput onChangeHandler={setFilterText} />
     </TableControlsWrapper>
@@ -363,7 +375,7 @@ export const DataTablesPane = ({
   return (
     <SouthPane data-test="some-purposeful-instance">
       <TabsWrapper contentHeight={tableSectionHeight}>
-        <CollapseWrapper>
+        <CollapseWrapper data-test="data-tab">
           <Collapse
             accordion
             bordered={false}
@@ -384,13 +396,27 @@ export const DataTablesPane = ({
                   tab={t('View results')}
                   key={RESULT_TYPES.results}
                 >
-                  {renderDataTable(RESULT_TYPES.results)}
+                  <DataTable
+                    isLoading={isLoading[RESULT_TYPES.results]}
+                    data={data[RESULT_TYPES.results]}
+                    columnNames={columnNames[RESULT_TYPES.results]}
+                    filterText={filterText}
+                    error={error[RESULT_TYPES.results]}
+                    errorMessage={errorMessage}
+                  />
                 </Tabs.TabPane>
                 <Tabs.TabPane
                   tab={t('View samples')}
                   key={RESULT_TYPES.samples}
                 >
-                  {renderDataTable(RESULT_TYPES.samples)}
+                  <DataTable
+                    isLoading={isLoading[RESULT_TYPES.samples]}
+                    data={data[RESULT_TYPES.samples]}
+                    columnNames={columnNames[RESULT_TYPES.samples]}
+                    filterText={filterText}
+                    error={error[RESULT_TYPES.samples]}
+                    errorMessage={errorMessage}
+                  />
                 </Tabs.TabPane>
               </Tabs>
             </Collapse.Panel>