From 29355577f148d1210c40043ef6028804469d2c30 Mon Sep 17 00:00:00 2001
From: "JUST.in DO IT" <justin.park@airbnb.com>
Date: Tue, 29 Aug 2023 14:38:07 -0700
Subject: [PATCH] fix(sqllab): error while removing a referenced table (#25114)

---
 .../SqlLab/components/TableElement/index.tsx  | 11 +++++--
 .../src/SqlLab/reducers/sqlLab.js             |  3 ++
 .../src/SqlLab/reducers/sqlLab.test.js        | 32 +++++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/TableElement/index.tsx b/superset-frontend/src/SqlLab/components/TableElement/index.tsx
index b003655634db3..979b9d85e9d5f 100644
--- a/superset-frontend/src/SqlLab/components/TableElement/index.tsx
+++ b/superset-frontend/src/SqlLab/components/TableElement/index.tsx
@@ -22,7 +22,7 @@ import type { Table } from 'src/SqlLab/types';
 import Collapse from 'src/components/Collapse';
 import Card from 'src/components/Card';
 import ButtonGroup from 'src/components/ButtonGroup';
-import { css, t, styled } from '@superset-ui/core';
+import { css, t, styled, useTheme } from '@superset-ui/core';
 import { debounce } from 'lodash';
 
 import {
@@ -102,6 +102,7 @@ const StyledCollapsePanel = styled(Collapse.Panel)`
 
 const TableElement = ({ table, ...props }: TableElementProps) => {
   const { dbId, schema, name, expanded } = table;
+  const theme = useTheme();
   const dispatch = useDispatch();
   const {
     data: tableMetadata,
@@ -253,7 +254,13 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
       );
     }
     return (
-      <ButtonGroup className="ws-el-controls">
+      <ButtonGroup
+        css={css`
+          display: flex;
+          column-gap: ${theme.gridUnit * 1.5}px;
+          margin-right: ${theme.gridUnit}px;
+        `}
+      >
         {keyLink}
         <IconTooltip
           className={
diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js
index 9f92372fd2fad..2b82f42d09b9c 100644
--- a/superset-frontend/src/SqlLab/reducers/sqlLab.js
+++ b/superset-frontend/src/SqlLab/reducers/sqlLab.js
@@ -184,6 +184,9 @@ export default function sqlLabReducer(state = {}, action) {
         if (action.query) {
           at.dataPreviewQueryId = action.query.id;
         }
+        if (existingTable.initialized) {
+          at.id = existingTable.id;
+        }
         return alterInArr(state, 'tables', existingTable, at);
       }
       // for new table, associate Id of query for data preview
diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.test.js b/superset-frontend/src/SqlLab/reducers/sqlLab.test.js
index 8827813698ea4..40597c41b0719 100644
--- a/superset-frontend/src/SqlLab/reducers/sqlLab.test.js
+++ b/superset-frontend/src/SqlLab/reducers/sqlLab.test.js
@@ -248,6 +248,38 @@ describe('sqlLabReducer', () => {
       expect(newState.tables).toHaveLength(1);
       expect(newState.tables[0].extra).toBe(true);
     });
+    it('should overwrite table ID be ignored when the existing table is already initialized', () => {
+      const action = {
+        type: actions.MERGE_TABLE,
+        table: newTable,
+      };
+      newState = sqlLabReducer(newState, action);
+      expect(newState.tables).toHaveLength(1);
+      // Merging the initialized remote id
+      const remoteId = 1;
+      const syncAction = {
+        type: actions.MERGE_TABLE,
+        table: {
+          ...newTable,
+          id: remoteId,
+          initialized: true,
+        },
+      };
+      newState = sqlLabReducer(newState, syncAction);
+      expect(newState.tables).toHaveLength(1);
+      expect(newState.tables[0].initialized).toBe(true);
+      expect(newState.tables[0].id).toBe(remoteId);
+      const overwriteAction = {
+        type: actions.MERGE_TABLE,
+        table: {
+          id: 'rnd_new_id',
+          ...newTable,
+        },
+      };
+      newState = sqlLabReducer(newState, overwriteAction);
+      expect(newState.tables).toHaveLength(1);
+      expect(newState.tables[0].id).toBe(remoteId);
+    });
     it('should expand and collapse a table', () => {
       const collapseTableAction = {
         type: actions.COLLAPSE_TABLE,