From c8eea2644795364906ab638fca24ce488927c498 Mon Sep 17 00:00:00 2001
From: justin-park <justin.park@airbnb.com>
Date: Fri, 2 Sep 2022 08:50:38 -0700
Subject: [PATCH] fix(sqllab): Copy link doesn't apply the unsaved changes

---
 .../ShareSqlLabQuery.test.jsx                 | 46 ++++++++++++++++++-
 .../components/ShareSqlLabQuery/index.tsx     | 22 ++++++---
 2 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/ShareSqlLabQuery.test.jsx b/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/ShareSqlLabQuery.test.jsx
index 8d87334e0fdb0..22913894fbc39 100644
--- a/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/ShareSqlLabQuery.test.jsx
+++ b/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/ShareSqlLabQuery.test.jsx
@@ -28,9 +28,10 @@ import '@testing-library/jest-dom/extend-expect';
 import userEvent from '@testing-library/user-event';
 import * as utils from 'src/utils/common';
 import ShareSqlLabQuery from 'src/SqlLab/components/ShareSqlLabQuery';
+import { initialState } from 'src/SqlLab/fixtures';
 
 const mockStore = configureStore([thunk]);
-const store = mockStore({});
+const store = mockStore(initialState);
 let isFeatureEnabledMock;
 
 const standardProvider = ({ children }) => (
@@ -41,6 +42,7 @@ const standardProvider = ({ children }) => (
 
 const defaultProps = {
   queryEditor: {
+    id: 'qe1',
     dbId: 0,
     name: 'query title',
     schema: 'query_schema',
@@ -51,6 +53,31 @@ const defaultProps = {
   addDangerToast: jest.fn(),
 };
 
+const unsavedQueryEditor = {
+  id: defaultProps.queryEditor.id,
+  dbId: 9888,
+  name: 'query title changed',
+  schema: 'query_schema_updated',
+  sql: 'SELECT * FROM Updated Limit 100',
+  autorun: true,
+};
+
+const standardProviderWithUnsaved = ({ children }) => (
+  <ThemeProvider theme={supersetTheme}>
+    <Provider
+      store={mockStore({
+        ...initialState,
+        sqlLab: {
+          ...initialState.sqlLab,
+          unsavedQueryEditor,
+        },
+      })}
+    >
+      {children}
+    </Provider>
+  </ThemeProvider>
+);
+
 describe('ShareSqlLabQuery', () => {
   const storeQueryUrl = 'glob:*/kv/store/';
   const storeQueryMockId = '123';
@@ -83,9 +110,26 @@ describe('ShareSqlLabQuery', () => {
         });
       });
       const button = screen.getByRole('button');
+      const { id, remoteId, ...expected } = defaultProps.queryEditor;
+      const storeQuerySpy = jest.spyOn(utils, 'storeQuery');
+      userEvent.click(button);
+      expect(storeQuerySpy.mock.calls).toHaveLength(1);
+      expect(storeQuerySpy).toBeCalledWith(expected);
+      storeQuerySpy.mockRestore();
+    });
+
+    it('calls storeQuery() with unsaved changes', async () => {
+      await act(async () => {
+        render(<ShareSqlLabQuery {...defaultProps} />, {
+          wrapper: standardProviderWithUnsaved,
+        });
+      });
+      const button = screen.getByRole('button');
+      const { id, ...expected } = unsavedQueryEditor;
       const storeQuerySpy = jest.spyOn(utils, 'storeQuery');
       userEvent.click(button);
       expect(storeQuerySpy.mock.calls).toHaveLength(1);
+      expect(storeQuerySpy).toBeCalledWith(expected);
       storeQuerySpy.mockRestore();
     });
   });
diff --git a/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/index.tsx b/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/index.tsx
index c521cb5dc75d9..37481bdee9249 100644
--- a/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/index.tsx
+++ b/superset-frontend/src/SqlLab/components/ShareSqlLabQuery/index.tsx
@@ -17,6 +17,7 @@
  * under the License.
  */
 import React from 'react';
+import { shallowEqual, useSelector } from 'react-redux';
 import { t, useTheme, styled } from '@superset-ui/core';
 import Button from 'src/components/Button';
 import Icons from 'src/components/Icons';
@@ -25,7 +26,7 @@ import CopyToClipboard from 'src/components/CopyToClipboard';
 import { storeQuery } from 'src/utils/common';
 import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import { FeatureFlag, isFeatureEnabled } from 'src/featureFlags';
-import { QueryEditor } from 'src/SqlLab/types';
+import { QueryEditor, SqlLabRootState } from 'src/SqlLab/types';
 
 interface ShareSqlLabQueryPropTypes {
   queryEditor: QueryEditor;
@@ -48,8 +49,18 @@ function ShareSqlLabQuery({
 }: ShareSqlLabQueryPropTypes) {
   const theme = useTheme();
 
+  const { dbId, name, schema, autorun, sql, remoteId } = useSelector<
+    SqlLabRootState,
+    Partial<QueryEditor>
+  >(({ sqlLab: { unsavedQueryEditor } }) => {
+    const { dbId, name, schema, autorun, sql, remoteId } = {
+      ...queryEditor,
+      ...(unsavedQueryEditor.id === queryEditor.id && unsavedQueryEditor),
+    };
+    return { dbId, name, schema, autorun, sql, remoteId };
+  }, shallowEqual);
+
   const getCopyUrlForKvStore = (callback: Function) => {
-    const { dbId, name, schema, autorun, sql } = queryEditor;
     const sharedQuery = { dbId, name, schema, autorun, sql };
 
     return storeQuery(sharedQuery)
@@ -66,10 +77,10 @@ function ShareSqlLabQuery({
   const getCopyUrlForSavedQuery = (callback: Function) => {
     let savedQueryToastContent;
 
-    if (queryEditor.remoteId) {
+    if (remoteId) {
       savedQueryToastContent = `${
         window.location.origin + window.location.pathname
-      }?savedQueryId=${queryEditor.remoteId}`;
+      }?savedQueryId=${remoteId}`;
       callback(savedQueryToastContent);
     } else {
       savedQueryToastContent = t('Please save the query to enable sharing');
@@ -101,8 +112,7 @@ function ShareSqlLabQuery({
   };
 
   const canShare =
-    !!queryEditor.remoteId ||
-    isFeatureEnabled(FeatureFlag.SHARE_QUERIES_VIA_KV_STORE);
+    !!remoteId || isFeatureEnabled(FeatureFlag.SHARE_QUERIES_VIA_KV_STORE);
 
   return (
     <>