fix: pass catalog when estimating query cost (#28410)

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 || '{}'),

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,

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"}
  )