From ba508a786c2a33486155ef03d3fb9fb44cb69ec9 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 26 Jul 2023 18:30:05 -0700
Subject: [PATCH] fix: pass schema on dataset creation (#24815)

---
 .../SaveDatasetModal.test.tsx                 | 42 ++++++++++++++++++-
 .../components/SaveDatasetModal/index.tsx     |  1 +
 superset-frontend/src/SqlLab/fixtures.ts      |  1 +
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx b/superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx
index f707998d1df67..4cac5c6204640 100644
--- a/superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx
+++ b/superset-frontend/src/SqlLab/components/SaveDatasetModal/SaveDatasetModal.test.tsx
@@ -18,10 +18,17 @@
  */
 import React from 'react';
 import * as reactRedux from 'react-redux';
-import { render, screen, cleanup, waitFor } from 'spec/helpers/testing-library';
+import {
+  fireEvent,
+  render,
+  screen,
+  cleanup,
+  waitFor,
+} from 'spec/helpers/testing-library';
 import userEvent from '@testing-library/user-event';
 import fetchMock from 'fetch-mock';
 import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
+import { createDatasource } from 'src/SqlLab/actions/sqlLab';
 import { user, testQuery, mockdatasets } from 'src/SqlLab/fixtures';
 
 const mockedProps = {
@@ -46,6 +53,15 @@ beforeEach(() => {
   cleanup();
 });
 
+// Mock the createDatasource action
+const useDispatchMock = jest.spyOn(reactRedux, 'useDispatch');
+jest.mock('src/SqlLab/actions/sqlLab', () => ({
+  createDatasource: jest.fn(),
+}));
+jest.mock('src/explore/exploreUtils/formData', () => ({
+  postFormData: jest.fn(),
+}));
+
 describe('SaveDatasetModal', () => {
   it('renders a "Save as new" field', () => {
     render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });
@@ -175,4 +191,28 @@ describe('SaveDatasetModal', () => {
     expect(screen.getByRole('button', { name: /back/i })).toBeVisible();
     expect(screen.getByRole('button', { name: /overwrite/i })).toBeVisible();
   });
+
+  it('sends the schema when creating the dataset', async () => {
+    const dummyDispatch = jest.fn().mockResolvedValue({});
+    useDispatchMock.mockReturnValue(dummyDispatch);
+    useSelectorMock.mockReturnValue({ ...user });
+
+    render(<SaveDatasetModal {...mockedProps} />, { useRedux: true });
+
+    const inputFieldText = screen.getByDisplayValue(/unimportant/i);
+    fireEvent.change(inputFieldText, { target: { value: 'my dataset' } });
+
+    const saveConfirmationBtn = screen.getByRole('button', {
+      name: /save/i,
+    });
+    userEvent.click(saveConfirmationBtn);
+
+    expect(createDatasource).toHaveBeenCalledWith({
+      datasourceName: 'my dataset',
+      dbId: 1,
+      schema: 'main',
+      sql: 'SELECT *',
+      templateParams: undefined,
+    });
+  });
 });
diff --git a/superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx b/superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx
index a42928608a9ab..7f605967ad9f6 100644
--- a/superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx
+++ b/superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx
@@ -296,6 +296,7 @@ export const SaveDatasetModal = ({
       createDatasource({
         sql: datasource.sql,
         dbId: datasource.dbId || datasource?.database?.id,
+        schema: datasource?.schema,
         templateParams,
         datasourceName: datasetName,
       }),
diff --git a/superset-frontend/src/SqlLab/fixtures.ts b/superset-frontend/src/SqlLab/fixtures.ts
index 658c1432dbb0a..0afd1c414977f 100644
--- a/superset-frontend/src/SqlLab/fixtures.ts
+++ b/superset-frontend/src/SqlLab/fixtures.ts
@@ -689,6 +689,7 @@ export const queryId = 'clientId2353';
 export const testQuery: ISaveableDatasource = {
   name: 'unimportant',
   dbId: 1,
+  schema: 'main',
   sql: 'SELECT *',
   columns: [
     {