From 498c6086670456b5a704310cd516b6624b4d6c84 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Thu, 9 May 2024 21:56:11 -0400
Subject: [PATCH] fix: pass catalog when estimating query cost (#28410)

---
 .../src/SqlLab/actions/sqlLab.js              |  3 +-
 .../components/SqlEditor/SqlEditor.test.tsx   | 82 +++++++++++++++++++
 superset/sqllab/schemas.py                    |  3 +
 3 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js
index a153c4eb4..7a73f9b9c 100644
--- a/superset-frontend/src/SqlLab/actions/sqlLab.js
+++ b/superset-frontend/src/SqlLab/actions/sqlLab.js
@@ -191,12 +191,13 @@ export function scheduleQuery(query) {
 
 export function estimateQueryCost(queryEditor) {
   return (dispatch, getState) => {
-    const { dbId, schema, sql, selectedText, templateParams } =
+    const { dbId, catalog, schema, sql, selectedText, templateParams } =
       getUpToDateQuery(getState(), queryEditor);
     const requestSql = selectedText || sql;
 
     const postPayload = {
       database_id: dbId,
+      catalog,
       schema,
       sql: requestSql,
       template_params: JSON.parse(templateParams || '{}'),
diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx
index f174110a2..a21b597c9 100644
--- a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx
+++ b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx
@@ -294,6 +294,88 @@ describe('SqlEditor', () => {
     ).toBeInTheDocument();
   });
 
+  describe('with EstimateQueryCost enabled', () => {
+    let isFeatureEnabledMock: jest.MockInstance<
+      boolean,
+      [feature: FeatureFlag]
+    >;
+    beforeEach(() => {
+      isFeatureEnabledMock = jest
+        .spyOn(uiCore, 'isFeatureEnabled')
+        .mockImplementation(
+          featureFlag => featureFlag === uiCore.FeatureFlag.EstimateQueryCost,
+        );
+    });
+    afterEach(() => {
+      isFeatureEnabledMock.mockClear();
+    });
+
+    it('sends the catalog and schema to the endpoint', async () => {
+      const estimateApi = 'http://localhost/api/v1/sqllab/estimate/';
+      fetchMock.post(estimateApi, {});
+
+      store = createStore({
+        ...initialState,
+        sqlLab: {
+          ...initialState.sqlLab,
+          databases: {
+            2023: {
+              allow_ctas: false,
+              allow_cvas: false,
+              allow_dml: false,
+              allow_file_upload: false,
+              allow_run_async: false,
+              backend: 'postgresql',
+              database_name: 'examples',
+              expose_in_sqllab: true,
+              force_ctas_schema: null,
+              id: 1,
+              allows_cost_estimate: true,
+            },
+          },
+          unsavedQueryEditor: {
+            id: defaultQueryEditor.id,
+            dbId: 2023,
+            sql: 'SELECT * FROM t',
+            schema: 'public',
+            catalog: 'prod',
+          },
+        },
+      });
+      const { findByText } = setup(mockedProps, store);
+      const button = await findByText('Estimate cost');
+      expect(button).toBeInTheDocument();
+
+      // click button
+      fireEvent.click(button);
+      await waitFor(() => {
+        expect(fetchMock.lastUrl()).toEqual(estimateApi);
+        expect(fetchMock.lastOptions()).toEqual(
+          expect.objectContaining({
+            body: JSON.stringify({
+              database_id: 2023,
+              catalog: 'prod',
+              schema: 'public',
+              sql: 'SELECT * FROM t',
+              template_params: {},
+            }),
+            cache: 'default',
+            credentials: 'same-origin',
+            headers: {
+              Accept: 'application/json',
+              'Content-Type': 'application/json',
+              'X-CSRFToken': '1234',
+            },
+            method: 'POST',
+            mode: 'same-origin',
+            redirect: 'follow',
+            signal: undefined,
+          }),
+        );
+      });
+    });
+  });
+
   describe('with SqllabBackendPersistence enabled', () => {
     let isFeatureEnabledMock: jest.MockInstance<
       boolean,
diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py
index 5e22a97e2..f441cdf04 100644
--- a/superset/sqllab/schemas.py
+++ b/superset/sqllab/schemas.py
@@ -37,6 +37,9 @@ class EstimateQueryCostSchema(Schema):
     template_params = fields.Dict(
         keys=fields.String(), metadata={"description": "The SQL query template params"}
     )
+    catalog = fields.String(
+        allow_none=True, metadata={"description": "The database catalog"}
+    )
     schema = fields.String(
         allow_none=True, metadata={"description": "The database schema"}
     )