From 9d5268ab6dbd6d93b9bb4047cddc99afc510d3c7 Mon Sep 17 00:00:00 2001
From: Usiel Riedl <usiel.riedl@automattic.com>
Date: Fri, 23 Aug 2024 08:07:59 +0800
Subject: [PATCH] feat(sqllab): Adds refresh button to table metadata in SQL
 Lab (#29974)

---
 .../TableElement/TableElement.test.tsx        | 37 ++++++++++++++++---
 .../SqlLab/components/TableElement/index.tsx  | 17 ++++++++-
 .../src/hooks/apiResources/tables.ts          |  7 ++++
 3 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx b/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx
index 30f41e5ee..2f391c269 100644
--- a/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx
+++ b/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx
@@ -51,11 +51,13 @@ const getTableMetadataEndpoint =
   /\/api\/v1\/database\/\d+\/table_metadata\/(?:\?.*)?$/;
 const getExtraTableMetadataEndpoint =
   /\/api\/v1\/database\/\d+\/table_metadata\/extra\/(?:\?.*)?$/;
-const updateTableSchemaEndpoint = 'glob:*/tableschemaview/*/expanded';
+const updateTableSchemaExpandedEndpoint = 'glob:*/tableschemaview/*/expanded';
+const updateTableSchemaEndpoint = 'glob:*/tableschemaview/';
 
 beforeEach(() => {
   fetchMock.get(getTableMetadataEndpoint, table);
   fetchMock.get(getExtraTableMetadataEndpoint, {});
+  fetchMock.post(updateTableSchemaExpandedEndpoint, {});
   fetchMock.post(updateTableSchemaEndpoint, {});
 });
 
@@ -84,7 +86,7 @@ test('has 4 IconTooltip elements', async () => {
     initialState,
   });
   await waitFor(() =>
-    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(4),
+    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(5),
   );
 });
 
@@ -104,7 +106,7 @@ test('fades table', async () => {
     initialState,
   });
   await waitFor(() =>
-    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(4),
+    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(5),
   );
   const style = window.getComputedStyle(getAllByTestId('fade')[0]);
   expect(style.opacity).toBe('0');
@@ -125,7 +127,7 @@ test('sorts columns', async () => {
     },
   );
   await waitFor(() =>
-    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(4),
+    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(5),
   );
   expect(
     getAllByTestId('mock-column-element').map(el => el.textContent),
@@ -154,7 +156,7 @@ test('removes the table', async () => {
     },
   );
   await waitFor(() =>
-    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(4),
+    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(5),
   );
   expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(0);
   fireEvent.click(getByText('Remove table preview'));
@@ -174,6 +176,29 @@ test('fetches table metadata when expanded', async () => {
   await waitFor(() =>
     expect(fetchMock.calls(getTableMetadataEndpoint)).toHaveLength(1),
   );
-  expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(0);
+  expect(fetchMock.calls(updateTableSchemaExpandedEndpoint)).toHaveLength(0);
   expect(fetchMock.calls(getExtraTableMetadataEndpoint)).toHaveLength(1);
 });
+
+test('refreshes table metadata when triggered', async () => {
+  const { getAllByTestId, getByText } = render(
+    <TableElement {...mockedProps} />,
+    {
+      useRedux: true,
+      initialState,
+    },
+  );
+  await waitFor(() =>
+    expect(getAllByTestId('mock-icon-tooltip')).toHaveLength(5),
+  );
+  expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(0);
+  expect(fetchMock.calls(getTableMetadataEndpoint)).toHaveLength(1);
+
+  fireEvent.click(getByText('Refresh table schema'));
+  await waitFor(() =>
+    expect(fetchMock.calls(getTableMetadataEndpoint)).toHaveLength(2),
+  );
+  await waitFor(() =>
+    expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(1),
+  );
+});
diff --git a/superset-frontend/src/SqlLab/components/TableElement/index.tsx b/superset-frontend/src/SqlLab/components/TableElement/index.tsx
index faa56db6b..824c9ec3c 100644
--- a/superset-frontend/src/SqlLab/components/TableElement/index.tsx
+++ b/superset-frontend/src/SqlLab/components/TableElement/index.tsx
@@ -32,6 +32,7 @@ import {
   syncTable,
 } from 'src/SqlLab/actions/sqlLab';
 import {
+  tableApiUtil,
   useTableExtendedMetadataQuery,
   useTableMetadataQuery,
 } from 'src/hooks/apiResources';
@@ -107,7 +108,7 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
   const {
     currentData: tableMetadata,
     isSuccess: isMetadataSuccess,
-    isLoading: isMetadataLoading,
+    isFetching: isMetadataFetching,
     isError: hasMetadataError,
   } = useTableMetadataQuery(
     {
@@ -177,6 +178,13 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
     setSortColumns(prevState => !prevState);
   };
 
+  const refreshTableMetadata = () => {
+    dispatch(
+      tableApiUtil.invalidateTags([{ type: 'TableMetadatas', id: name }]),
+    );
+    dispatch(syncTable(table, tableData));
+  };
+
   const renderWell = () => {
     let partitions;
     let metadata;
@@ -268,6 +276,11 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
           }
         `}
       >
+        <IconTooltip
+          className="fa fa-refresh pull-left m-l-2 pointer"
+          onClick={refreshTableMetadata}
+          tooltip={t('Refresh table schema')}
+        />
         {keyLink}
         <IconTooltip
           className={
@@ -341,7 +354,7 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
         </Tooltip>
 
         <div className="pull-right header-right-side">
-          {isMetadataLoading || isExtraMetadataLoading ? (
+          {isMetadataFetching || isExtraMetadataLoading ? (
             <Loading position="inline" />
           ) : (
             <Fade
diff --git a/superset-frontend/src/hooks/apiResources/tables.ts b/superset-frontend/src/hooks/apiResources/tables.ts
index f2accf499..86b080745 100644
--- a/superset-frontend/src/hooks/apiResources/tables.ts
+++ b/superset-frontend/src/hooks/apiResources/tables.ts
@@ -117,6 +117,13 @@ const tableApi = api.injectEndpoints({
       }),
     }),
     tableMetadata: builder.query<TableMetaData, FetchTableMetadataQueryParams>({
+      providesTags: result =>
+        result
+          ? [
+              { type: 'TableMetadatas', id: result.name },
+              { type: 'TableMetadatas', id: 'LIST' },
+            ]
+          : [{ type: 'TableMetadatas', id: 'LIST' }],
       query: ({ dbId, catalog, schema, table }) => ({
         endpoint: `/api/v1/database/${dbId}/table_metadata/${toQueryString({
           name: table,