From 874ee59a4a88d324f8223f9943c08ddaa52de3af Mon Sep 17 00:00:00 2001 From: elizabeth-ilina Date: Sun, 4 Oct 2026 20:50:54 -0400 Subject: [PATCH 01/58] Add some initial implementations for POC and proposal. --- docker/pythonpath_dev/superset_config.py | 1 + docs/docs/using-superset/exploring-data.mdx | 25 +++ .../components/SqlEditor/SqlEditor.test.tsx | 167 ++++++++++++++++++ .../src/SqlLab/components/SqlEditor/index.tsx | 125 ++++++++++++- .../DatabaseModal/ExtraOptions.test.tsx | 82 +++++++++ .../databases/DatabaseModal/ExtraOptions.tsx | 117 ++++++++++++ .../databases/DatabaseModal/index.test.tsx | 37 ++++ .../databases/DatabaseModal/index.tsx | 28 +++ .../src/features/databases/types.ts | 3 + superset/commands/sql_lab/estimate.py | 22 ++- superset/databases/api.py | 1 + superset/databases/schemas.py | 25 +++ superset/db_engine_specs/bigquery.py | 13 +- superset/models/core.py | 18 ++ superset/sqllab/schemas.py | 2 + superset/sqllab/utils.py | 1 + .../commands/sql_lab/test_estimate.py | 46 +++++ tests/unit_tests/databases/schema_tests.py | 43 +++++ .../db_engine_specs/test_bigquery.py | 13 ++ tests/unit_tests/models/core_test.py | 54 ++++++ tests/unit_tests/sqllab/test_utils.py | 20 +++ 21 files changed, 828 insertions(+), 15 deletions(-) diff --git a/docker/pythonpath_dev/superset_config.py b/docker/pythonpath_dev/superset_config.py index a4553e00997e..6dfcd12e9d74 100644 --- a/docker/pythonpath_dev/superset_config.py +++ b/docker/pythonpath_dev/superset_config.py @@ -116,6 +116,7 @@ class CeleryConfig: FEATURE_FLAGS = { "ALERT_REPORTS": True, "DATASET_FOLDERS": True, + "ESTIMATE_QUERY_COST": True, "ENABLE_EXTENSIONS": True, "MOBILE_CONSUMPTION_MODE": True, "SEMANTIC_LAYERS": True, diff --git a/docs/docs/using-superset/exploring-data.mdx b/docs/docs/using-superset/exploring-data.mdx index 444f0ae5bd1f..d3454eca189f 100644 --- a/docs/docs/using-superset/exploring-data.mdx +++ b/docs/docs/using-superset/exploring-data.mdx @@ -409,6 +409,31 @@ to be drawn as zero across the entire day rather than only around the hours that Lastly, save your chart as Tutorial Resample and add it to the Tutorial Dashboard. Go to the tutorial dashboard to see the four charts side by side and compare the different outputs. +### BigQuery pre-run cost warnings + +Operators can enable `ESTIMATE_QUERY_COST` and turn on **Enable query cost +estimation** in a BigQuery connection's SQL Lab settings. Check **Enable pre-run +cost warnings**, then enter an amount and select **Bytes**, **MB**, or **GB**. +The amount is stored as integer bytes in `cost_estimate_warning_threshold_bytes`. +Unchecking the warning checkbox disables warnings and preserves the threshold; +clearing the amount disables warnings until you enter another threshold. +The checkbox is stored as `cost_estimate_warning_enabled` and requires query cost +estimation to be enabled. + +The SQL editor's Run button and keyboard shortcuts request a dry-run estimate +before executing. A positive estimate at or below the threshold proceeds +automatically; a larger estimate asks for confirmation. Failed, missing, zero-byte, +and unsupported-statement estimates also ask for confirmation. Changing the query +or its parameters while the estimate is pending discards the response; click Run +again to estimate the changed query. + +This advisory check supports single BigQuery SELECT statements, including selected +text. It does not cover CTAS/CVAS, query-history reruns, dashboards, scheduled +queries, or direct execution API calls. Bytes processed are an estimate, not an +exact bill or an enforced spending limit. Estimation and execution use separate +requests, so dynamic templates, SQL transformations, and changes in warehouse +data can affect the actual work. Warehouse billing limits remain applicable. + ### SQL Lab Tips **Schema and table browser**: The left-side table browser uses a collapsible treeview — click a schema to expand its tables, and click a table to see its columns and sample data inline. This makes navigating large schemas much faster than the previous flat list. diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx index d2bd1380ad4c..9c7077166a1b 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.test.tsx @@ -42,6 +42,7 @@ import { import ResultSet from 'src/SqlLab/components/ResultSet'; import { api } from 'src/hooks/apiResources/queryApi'; import setupCodeOverrides from 'src/setup/setupCodeOverrides'; +import { queryEditorSetSql } from 'src/SqlLab/actions/sqlLab'; import type { Action, Middleware, Store } from 'redux'; import SqlEditor, { Props } from '.'; @@ -205,6 +206,172 @@ describe('SqlEditor', () => { ).toBeInTheDocument(); }); + test.each([ + { bytes: [99], proceed: false, executions: 1, confirmation: false }, + { bytes: [100], proceed: false, executions: 1, confirmation: false }, + { bytes: [101], proceed: false, executions: 0, confirmation: true }, + { bytes: [101], proceed: true, executions: 1, confirmation: true }, + { bytes: [null], proceed: false, executions: 0, confirmation: true }, + { bytes: [null], proceed: true, executions: 1, confirmation: true }, + { bytes: [0], proceed: false, executions: 0, confirmation: true }, + { bytes: [50, 200], proceed: false, executions: 0, confirmation: true }, + ])( + 'pre-run estimate $bytes, confirmation $proceed submits $executions queries', + async ({ bytes, proceed, executions, confirmation }) => { + mockIsFeatureEnabled.mockImplementation( + flag => flag === FeatureFlag.EstimateQueryCost, + ); + fetchMock.post( + 'glob:*/api/v1/sqllab/estimate/', + { result: { estimate: [], bytes_processed: bytes } }, + { name: 'pre-run-estimate' }, + ); + store = createStore({ + ...mockInitialState, + sqlLab: { + ...mockInitialState.sqlLab, + databases: { + 1991: { + ...mockInitialState.sqlLab.databases[1991], + id: 1991, + backend: 'bigquery', + allows_cost_estimate: true, + cost_estimate_warning_threshold_bytes: 100, + }, + }, + }, + }); + try { + const { findByTestId, findByRole, queryByRole } = setup( + mockedProps, + store, + ); + fireEvent.click(await findByTestId('run-query-action')); + await waitFor(() => { + expect(fetchMock.callHistory.called('pre-run-estimate')).toBe(true); + }); + await act(async () => { + await fetchMock.callHistory.flush(true); + }); + if (confirmation) { + const decision = await findByRole('button', { + name: proceed ? 'Run anyway' : 'Cancel', + }); + await act(async () => { + fireEvent.click(decision); + }); + } else { + expect(queryByRole('dialog')).not.toBeInTheDocument(); + } + expect( + actions.filter(action => action.type === 'START_QUERY'), + ).toHaveLength(executions); + } finally { + fetchMock.removeRoute('pre-run-estimate'); + mockIsFeatureEnabled.mockReset(); + } + }, + ); + + test.each([false, true])( + 'pending estimate deduplicates clicks and discards changed SQL (%s)', + async changed => { + mockIsFeatureEnabled.mockImplementation( + flag => flag === FeatureFlag.EstimateQueryCost, + ); + let finishEstimate: (value: { + result: { bytes_processed: number[] }; + }) => void = () => {}; + const response = new Promise<{ result: { bytes_processed: number[] } }>( + resolve => { + finishEstimate = resolve; + }, + ); + fetchMock.post('glob:*/api/v1/sqllab/estimate/', () => response, { + name: 'pending-estimate', + }); + store = createStore({ + ...mockInitialState, + sqlLab: { + ...mockInitialState.sqlLab, + databases: { + 1991: { + ...mockInitialState.sqlLab.databases[1991], + id: 1991, + backend: 'bigquery', + allows_cost_estimate: true, + cost_estimate_warning_threshold_bytes: 100, + }, + }, + unsavedQueryEditor: { + ...mockInitialState.sqlLab.unsavedQueryEditor, + sql: 'SELECT 1', + catalog: 'project', + schema: 'dataset', + templateParams: '{"limit": 5}', + }, + }, + }); + try { + const { findByTestId, findByRole, queryByRole } = setup( + mockedProps, + store, + ); + const button = await findByTestId('run-query-action'); + fireEvent.click(button); + fireEvent.click(button); + await waitFor(() => { + expect(fetchMock.callHistory.calls('pending-estimate')).toHaveLength( + 1, + ); + }); + const body = JSON.parse( + String( + fetchMock.callHistory.calls('pending-estimate')[0].options.body, + ), + ); + expect(body).toEqual({ + database_id: 1991, + catalog: 'project', + schema: 'dataset', + sql: 'SELECT 1', + template_params: { limit: 5 }, + include_raw: true, + }); + expect( + actions.filter(action => action.type === 'START_QUERY'), + ).toHaveLength(0); + if (changed) { + act(() => { + store.dispatch( + queryEditorSetSql(mockedProps.queryEditor, 'SELECT 2'), + ); + }); + } + await act(async () => { + finishEstimate({ result: { bytes_processed: [101] } }); + await fetchMock.callHistory.flush(true); + }); + if (!changed) { + const continueButton = await findByRole('button', { + name: 'Run anyway', + }); + await act(async () => { + fireEvent.click(continueButton); + }); + } else { + expect(queryByRole('dialog')).not.toBeInTheDocument(); + } + expect( + actions.filter(action => action.type === 'START_QUERY'), + ).toHaveLength(changed ? 0 : 1); + } finally { + fetchMock.removeRoute('pending-estimate'); + mockIsFeatureEnabled.mockReset(); + } + }, + ); + test('renders db unavailable message', async () => { const queryEditor = initialState.sqlLab.queryEditors[1]; const { findByText } = setup({ ...mockedProps, queryEditor }, store); diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx index b480bffdec30..302438efcc60 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx @@ -30,7 +30,7 @@ import { import type { editors } from '@apache-superset/core'; import useEffectEvent from 'src/hooks/useEffectEvent'; -import { shallowEqual, useSelector } from 'react-redux'; +import { shallowEqual, useSelector, useStore } from 'react-redux'; import { useAppDispatch } from 'src/SqlLab/hooks/useAppDispatch'; import AutoSizer from 'react-virtualized-auto-sizer'; import { t } from '@apache-superset/core/translation'; @@ -41,6 +41,7 @@ import { QueryResponse, Query, QueryState, + SupersetClient, } from '@superset-ui/core'; import { Alert } from '@apache-superset/core/components'; import { css, styled, useTheme } from '@apache-superset/core/theme'; @@ -68,6 +69,7 @@ import { addNewQueryEditor, CtasEnum, estimateQueryCost, + getUpToDateQuery, persistEditorHeight, postStopQuery, queryEditorSetAutorun, @@ -241,6 +243,16 @@ const SqlEditor: FC = ({ }) => { const theme = useTheme(); const dispatch = useAppDispatch(); + const store = useStore(); + const [costWarningModal, costWarningContext] = Modal.useModal(); + const estimatePending = useRef(false); + const mounted = useRef(true); + useEffect(() => { + mounted.current = true; + return () => { + mounted.current = false; + }; + }, []); const { database, @@ -299,19 +311,18 @@ const SqlEditor: FC = ({ const successful = latestQuery?.state === QueryState.Success; const resultColumns = latestQuery?.results?.columns || []; - const startQuery = useCallback( + const executeQuery = useCallback( ( ctasArg = false, ctas_method: (typeof CtasEnum)[keyof typeof CtasEnum] = CtasEnum.Table, ) => { - if (!database) { - return; - } - + if (!database) return; dispatch( runQueryFromSqlEditor( - // eslint-disable-next-line @typescript-eslint/no-explicit-any - database as any, + { + id: database.id, + allow_run_async: Boolean(database.allow_run_async), + }, queryEditor, defaultQueryLimit, ctasArg ? ctas : '', @@ -324,6 +335,103 @@ const SqlEditor: FC = ({ [ctas, database, defaultQueryLimit, dispatch, queryEditor], ); + const startQuery = useCallback( + async ( + ctasArg = false, + ctas_method: (typeof CtasEnum)[keyof typeof CtasEnum] = CtasEnum.Table, + ) => { + if (!database || estimatePending.current) { + return; + } + + const warningThreshold = database.cost_estimate_warning_threshold_bytes; + if ( + !ctasArg && + isFeatureEnabled(FeatureFlag.EstimateQueryCost) && + database.backend === 'bigquery' && + database.allows_cost_estimate && + typeof warningThreshold === 'number' && + Number.isSafeInteger(warningThreshold) && + warningThreshold >= 0 + ) { + const getInputs = () => { + const qe = getUpToDateQuery(store.getState(), queryEditor); + return { + database_id: qe.dbId, + catalog: qe.catalog, + schema: qe.schema, + sql: qe.selectedText || qe.sql, + template_params: qe.templateParams || '{}', + query_limit: qe.queryLimit, + }; + }; + const inputs = getInputs(); + const isCurrent = () => + mounted.current && + JSON.stringify(inputs) === JSON.stringify(getInputs()); + estimatePending.current = true; + let estimatedBytes: unknown; + try { + const { json } = await SupersetClient.post({ + endpoint: '/api/v1/sqllab/estimate/', + jsonPayload: { + database_id: inputs.database_id, + catalog: inputs.catalog, + schema: inputs.schema, + sql: inputs.sql, + template_params: JSON.parse(inputs.template_params), + include_raw: true, + }, + }); + const bytes: unknown = json.result?.bytes_processed; + estimatedBytes = + Array.isArray(bytes) && bytes.length === 1 ? bytes[0] : undefined; + } catch { + estimatedBytes = undefined; + } + try { + if (!isCurrent()) return; + let message = t( + 'The query size could not be reliably estimated. Zero-byte estimates and unsupported statements may still process data.', + ); + if ( + typeof estimatedBytes === 'number' && + Number.isSafeInteger(estimatedBytes) && + estimatedBytes > 0 + ) { + if (estimatedBytes <= warningThreshold) { + executeQuery(ctasArg, ctas_method); + return; + } + message = t( + 'This query is estimated to process %(estimate)s bytes. The warning threshold is %(threshold)s bytes.', + { + estimate: estimatedBytes.toLocaleString(), + threshold: warningThreshold.toLocaleString(), + }, + ); + } + const confirmed = await costWarningModal.confirm({ + title: t('Query cost warning'), + content: message, + okText: t('Run anyway'), + cancelText: t('Cancel'), + autoFocusButton: 'cancel', + }); + if (confirmed && isCurrent()) { + executeQuery(ctasArg, ctas_method); + } + } finally { + estimatePending.current = false; + } + return; + } + + executeQuery(ctasArg, ctas_method); + }, + [costWarningModal, database, executeQuery, queryEditor, store], + ); + const formatCurrentQuery = useCallback( (useShortcut?: boolean) => { logAction(LOG_ACTIONS_SQLLAB_FORMAT_SQL, { @@ -1095,6 +1203,7 @@ const SqlEditor: FC = ({ {t('Name')} + {costWarningContext} ); }; diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx index abc2a6063716..4262b50e786f 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx @@ -21,6 +21,7 @@ import { fireEvent, render, screen, + selectOption, waitFor, } from 'spec/helpers/testing-library'; import { t } from '@apache-superset/core/translation'; @@ -120,6 +121,87 @@ describe('ExtraOptions Component', () => { expect(onExtraInputChange).toHaveBeenCalled(); }); + test.each([false, true])( + 'pre-run warnings require estimation (%s)', + estimationEnabled => { + renderComponent({ + ...defaultDb, + backend: 'bigquery', + extra: JSON.stringify({ cost_estimate_enabled: estimationEnabled }), + }); + fireEvent.click(screen.getByText('SQL Lab')); + const checkbox = screen.getByRole('checkbox', { + name: 'Enable pre-run cost warnings', + }); + expect(checkbox).not.toBeChecked(); + if (estimationEnabled) { + expect(checkbox).toBeEnabled(); + fireEvent.click(checkbox); + expect(onExtraInputChange).toHaveBeenCalledWith( + expect.objectContaining({ + target: expect.objectContaining({ + name: 'cost_estimate_warning_enabled', + checked: true, + }), + }), + ); + } else { + expect(checkbox).toBeDisabled(); + } + expect( + screen.queryByLabelText('Warn when estimated data processed exceeds'), + ).not.toBeInTheDocument(); + }, + ); + + test('warning amount and size selector convert MB and GB to bytes', async () => { + renderComponent({ + ...defaultDb, + backend: 'bigquery', + extra: JSON.stringify({ + cost_estimate_enabled: true, + cost_estimate_warning_enabled: true, + cost_estimate_warning_threshold_bytes: 100000000, + }), + }); + fireEvent.click(screen.getByText('SQL Lab')); + expect( + screen.getByRole('checkbox', { name: 'Enable pre-run cost warnings' }), + ).toBeChecked(); + const amount = screen.getByLabelText( + 'Warn when estimated data processed exceeds', + ); + expect(amount).toHaveValue(100); + fireEvent.change(amount, { target: { value: '250' } }); + expect(onExtraInputChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + target: expect.objectContaining({ + name: 'cost_estimate_warning_threshold_bytes', + value: '250000000', + }), + }), + ); + await selectOption('GB', 'Threshold size unit'); + expect(onExtraInputChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + target: expect.objectContaining({ value: '100000000000' }), + }), + ); + fireEvent.change(amount, { target: { value: '0.5' } }); + expect(onExtraInputChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + target: expect.objectContaining({ value: '500000000' }), + }), + ); + await selectOption('Bytes', 'Threshold size unit'); + fireEvent.change(amount, { target: { value: '123' } }); + expect(onExtraInputChange).toHaveBeenLastCalledWith( + expect.objectContaining({ + target: expect.objectContaining({ value: '123' }), + }), + ); + }); + test('calls onExtraEditorChange when metadata_params json editor changes', async () => { renderComponent(); diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx index cb5931bc58db..08eecd4c8901 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx @@ -31,6 +31,8 @@ import { Collapse, InfoTooltip, CollapseLabelInModal, + Flex, + Select, type CheckboxChangeEvent, } from '@superset-ui/core/components'; import { useJsonValidation } from '@superset-ui/core/components/AsyncAceEditor'; @@ -42,6 +44,9 @@ import { } from './styles'; import { DatabaseObject, ExtraJson } from '../types'; +const costWarningUnits = { Bytes: 1, MB: 1000000, GB: 1000000000 }; +type CostWarningUnit = keyof typeof costWarningUnits; + const ExtraOptions = ({ db, onInputChange, @@ -101,6 +106,29 @@ const ExtraOptions = ({ return value; }); + const warningBytes = extraJson.cost_estimate_warning_threshold_bytes; + const warningsEnabled = + extraJson.cost_estimate_warning_enabled ?? warningBytes !== undefined; + const [warningUnit, setWarningUnit] = useState(() => { + if (warningBytes === undefined) return 'MB'; + if (warningBytes >= 1e9 && warningBytes % 1e9 === 0) return 'GB'; + if (warningBytes >= 1e6 && warningBytes % 1e6 === 0) return 'MB'; + return 'Bytes'; + }); + const warningAmount = + warningBytes === undefined + ? '' + : warningBytes / costWarningUnits[warningUnit]; + const updateWarningBytes = (value: number | '') => { + onExtraInputChange({ + target: { + name: 'cost_estimate_warning_threshold_bytes', + type: 'number', + value: value === '' ? '' : String(Math.round(value)), + }, + } as ChangeEvent); + }; + // JSON validation hooks for the three editors const secureExtraAnnotations = useJsonValidation(db?.masked_encrypted_extra, { errorPrefix: 'Invalid secure extra JSON', @@ -288,6 +316,95 @@ const ExtraOptions = ({ /> + {(db?.backend === 'bigquery' || + db?.engine === 'bigquery') && ( + +
+ + {t('Enable pre-run cost warnings')} + +
+ {warningsEnabled && ( + + + + + updateWarningBytes( + event.target.value === '' + ? '' + : Number(event.target.value) * + costWarningUnits[warningUnit], + ) + } + /> + - updateWarningBytes( - event.target.value === '' - ? '' - : Number(event.target.value) * - costWarningUnits[warningUnit], - ) - } - /> - - - - {t( - 'The analyst can still choose to run the query. 1 MB = 1,000,000 bytes; 1 GB = 1,000,000,000 bytes.', - )} - - - )} - - )} + value={warningAmount} + disabled={!estimationEnabled} + onChange={onExtraInputChange} + /> + + {warningMetric === 'bytes_processed' + ? t( + 'Estimated bytes processed are usage estimates, not a guaranteed bill. You can still choose to run the query.', + ) + : t( + 'PostgreSQL planner cost uses planner units, not runtime or money. You can still choose to run the query.', + )} + + + )} + + )}
{ }); }); - test.each(['0', '1025', ''])( - 'cost warning threshold %s is saved as integer bytes or removed', - value => { - const currentState = dbReducer(databaseFixture, { + test('enabling warnings sets the metric without inventing a threshold', () => { + const state = dbReducer( + { + ...databaseFixture, + engine_information: { query_cost_warning_metric: 'total_cost' }, + }, + { type: ActionType.ExtraInputChange, - payload: { name: 'cost_estimate_warning_threshold_bytes', value }, - }); - expect(JSON.parse(currentState.extra || '{}')).toEqual( - value === '' - ? { cost_estimate_warning_enabled: false } - : { - cost_estimate_warning_enabled: true, - cost_estimate_warning_threshold_bytes: Number(value), - }, + payload: { name: 'query_cost_warning_enabled', checked: true }, + }, + ); + expect(JSON.parse(state?.extra || '{}')).toEqual({ + query_cost_warning: { enabled: true, metric: 'total_cost' }, + }); + }); + + test.each([100, 'invalid'])( + 'disabling a warning preserves valid values and removes invalid ones (%s)', + threshold => { + const extra = { + unrelated: { keep: true }, + cost_estimate_warning_threshold_bytes: 55, + query_cost_warning: { + enabled: true, + metric: 'bytes_processed', + threshold, + }, + }; + const state = dbReducer( + { ...databaseFixture, extra: JSON.stringify(extra) }, + { + type: ActionType.ExtraInputChange, + payload: { name: 'query_cost_warning_enabled', checked: false }, + }, ); + expect(JSON.parse(state?.extra || '{}')).toEqual({ + ...extra, + query_cost_warning: { + enabled: false, + metric: 'bytes_processed', + ...(typeof threshold === 'number' && { threshold }), + }, + }); }, ); - test('cost warnings can be disabled without losing the threshold', () => { - const enabled = dbReducer(databaseFixture, { - type: ActionType.ExtraInputChange, - payload: { name: 'cost_estimate_warning_enabled', checked: true }, - }); - expect(JSON.parse(enabled.extra || '{}')).toEqual({ - cost_estimate_warning_enabled: true, - cost_estimate_warning_threshold_bytes: 100000000, - }); - const disabled = dbReducer(enabled, { - type: ActionType.ExtraInputChange, - payload: { name: 'cost_estimate_warning_enabled', checked: false }, + test('editing a PostgreSQL threshold keeps decimals and unrelated settings', () => { + const state = dbReducer( + { + ...databaseFixture, + engine_information: { query_cost_warning_metric: 'total_cost' }, + extra: JSON.stringify({ + unrelated: true, + query_cost_warning: { + enabled: true, + metric: 'total_cost', + threshold: 10, + }, + }), + }, + { + type: ActionType.ExtraInputChange, + payload: { name: 'query_cost_warning_threshold', value: '12.5' }, + }, + ); + expect(JSON.parse(state?.extra || '{}')).toEqual({ + unrelated: true, + query_cost_warning: { + enabled: true, + metric: 'total_cost', + threshold: 12.5, + }, }); - expect(JSON.parse(disabled.extra || '{}')).toEqual({ - cost_estimate_warning_enabled: false, - cost_estimate_warning_threshold_bytes: 100000000, + }); + + test('turning estimation off retains the warning rule', () => { + const extra = { + query_cost_warning: { + enabled: true, + metric: 'total_cost', + threshold: 12.5, + }, + cost_estimate_enabled: true, + }; + const state = dbReducer( + { ...databaseFixture, extra: JSON.stringify(extra) }, + { + type: ActionType.ExtraInputChange, + payload: { + name: 'cost_estimate_enabled', + type: 'checkbox', + checked: false, + }, + }, + ); + expect(JSON.parse(state?.extra || '{}')).toEqual({ + ...extra, + cost_estimate_enabled: false, }); }); @@ -2834,3 +2898,26 @@ describe('dbReducer', () => { }); }); }); + +test('changing a warning metric does not reinterpret an old threshold', () => { + const state = dbReducer( + { + ...databaseFixture, + engine_information: { query_cost_warning_metric: 'total_cost' }, + extra: JSON.stringify({ + query_cost_warning: { + enabled: false, + metric: 'bytes_processed', + threshold: 100, + }, + }), + }, + { + type: ActionType.ExtraInputChange, + payload: { name: 'query_cost_warning_enabled', checked: true }, + }, + ); + expect(JSON.parse(state?.extra || '{}')).toEqual({ + query_cost_warning: { enabled: true, metric: 'total_cost' }, + }); +}); diff --git a/superset-frontend/src/features/databases/DatabaseModal/index.tsx b/superset-frontend/src/features/databases/DatabaseModal/index.tsx index b2ab5f06ff05..e5160f162415 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/index.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/index.tsx @@ -77,6 +77,10 @@ import { DatabaseParameters, } from '../types'; import ExtraOptions from './ExtraOptions'; +import { + disableQueryCostWarning, + getQueryCostWarning, +} from './queryCostWarning'; import SqlAlchemyForm from './SqlAlchemyForm'; import DatabaseConnectionForm from './DatabaseConnectionForm'; import { @@ -339,31 +343,39 @@ export function dbReducer( }; } case ActionType.ExtraInputChange: - if (action.payload.name === 'cost_estimate_warning_enabled') { - return { - ...trimmedState, - extra: JSON.stringify({ - ...extraJson, - cost_estimate_warning_enabled: action.payload.checked, - cost_estimate_warning_threshold_bytes: - extraJson.cost_estimate_warning_threshold_bytes ?? - (action.payload.checked ? 100000000 : undefined), - }), - }; - } - if (action.payload.name === 'cost_estimate_warning_threshold_bytes') { + if ( + action.payload.name === 'query_cost_warning_enabled' || + action.payload.name === 'query_cost_warning_threshold' + ) { + const rule = getQueryCostWarning( + extraJson, + trimmedState.backend || trimmedState.engine, + ); + const metric = + trimmedState.engine_information?.query_cost_warning_metric; + const queryCostWarning = + action.payload.name === 'query_cost_warning_enabled' + ? action.payload.checked + ? { + ...rule, + enabled: true, + metric, + ...(rule.metric !== metric && { threshold: undefined }), + } + : disableQueryCostWarning(rule) + : { + ...rule, + metric, + threshold: + action.payload.value === '' + ? undefined + : Number(action.payload.value), + }; return { ...trimmedState, extra: JSON.stringify({ ...extraJson, - cost_estimate_warning_enabled: - action.payload.value === '' - ? false - : (extraJson.cost_estimate_warning_enabled ?? true), - cost_estimate_warning_threshold_bytes: - action.payload.value === '' - ? undefined - : Number(action.payload.value), + query_cost_warning: queryCostWarning, }), }; } diff --git a/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.test.ts b/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.test.ts new file mode 100644 index 000000000000..22c1a7989d33 --- /dev/null +++ b/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.test.ts @@ -0,0 +1,94 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { + disableQueryCostWarning, + getQueryCostWarning, +} from './queryCostWarning'; +import { ExtraJson } from '../types'; + +test.each([undefined, false, true, 'true', 1])( + 'legacy opt-in must be exactly true (%s)', + enabled => { + const extra = { + cost_estimate_warning_enabled: enabled, + cost_estimate_warning_threshold_bytes: 100, + } as ExtraJson; + expect(getQueryCostWarning(extra, 'bigquery')).toEqual({ + enabled: enabled === true, + metric: 'bytes_processed', + threshold: 100, + }); + }, +); + +test.each([-1, 0, 1.5, '100', true, null, 2 ** 53])( + 'invalid legacy threshold stays off without a default (%s)', + threshold => { + expect( + getQueryCostWarning( + { + cost_estimate_warning_enabled: true, + cost_estimate_warning_threshold_bytes: threshold, + } as ExtraJson, + 'bigquery', + ), + ).toEqual({ enabled: false }); + }, +); + +test('legacy bytes do not become a PostgreSQL planner threshold', () => { + expect( + getQueryCostWarning( + { + cost_estimate_warning_enabled: true, + cost_estimate_warning_threshold_bytes: 100, + }, + 'postgresql', + ), + ).toEqual({ enabled: false }); +}); + +test.each([null, { enabled: false }, { enabled: true, threshold: 'invalid' }])( + 'an explicit new rule prevents legacy fallback (%j)', + rule => { + expect( + getQueryCostWarning( + { + query_cost_warning: rule, + cost_estimate_warning_enabled: true, + cost_estimate_warning_threshold_bytes: 100, + } as ExtraJson, + 'bigquery', + ), + ).toEqual(rule || {}); + }, +); + +test.each([NaN, Infinity, 0, -1, '100', true, 0.5])( + 'disabling drops invalid byte thresholds (%s)', + threshold => { + expect( + disableQueryCostWarning({ + enabled: true, + metric: 'bytes_processed', + threshold, + }), + ).toEqual({ enabled: false, metric: 'bytes_processed' }); + }, +); diff --git a/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.ts b/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.ts new file mode 100644 index 000000000000..c58603b53028 --- /dev/null +++ b/superset-frontend/src/features/databases/DatabaseModal/queryCostWarning.ts @@ -0,0 +1,68 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { ExtraJson, QueryCostWarning } from '../types'; + +/** Read the explicit rule or a valid BigQuery draft without implicit opt-in. */ +export function getQueryCostWarning( + extra: ExtraJson, + engine?: string, +): Record { + if ('query_cost_warning' in extra) { + const rule: unknown = extra.query_cost_warning; + return rule && typeof rule === 'object' && !Array.isArray(rule) + ? { ...(rule as Record) } + : {}; + } + const threshold = extra.cost_estimate_warning_threshold_bytes; + if ( + engine === 'bigquery' && + typeof threshold === 'number' && + Number.isSafeInteger(threshold) && + threshold > 0 + ) { + return { + enabled: extra.cost_estimate_warning_enabled === true, + metric: 'bytes_processed', + threshold, + }; + } + return { enabled: false }; +} + +/** Disable warnings while retaining only usable saved metric and threshold values. */ +export function disableQueryCostWarning( + rule: Record, +): QueryCostWarning { + const metric = + rule.metric === 'bytes_processed' || rule.metric === 'total_cost' + ? rule.metric + : undefined; + const threshold = + typeof rule.threshold === 'number' && + Number.isFinite(rule.threshold) && + rule.threshold > 0 && + (metric !== 'bytes_processed' || Number.isSafeInteger(rule.threshold)) + ? rule.threshold + : undefined; + return { + enabled: false, + ...(metric && { metric }), + ...(threshold !== undefined && { threshold }), + }; +} diff --git a/superset-frontend/src/features/databases/types.ts b/superset-frontend/src/features/databases/types.ts index d2b60d153027..2e91d6185d30 100644 --- a/superset-frontend/src/features/databases/types.ts +++ b/superset-frontend/src/features/databases/types.ts @@ -64,6 +64,14 @@ export type DatabaseParameters = { project_id?: string; }; +export type QueryCostWarningMetric = 'bytes_processed' | 'total_cost'; + +export type QueryCostWarning = { + enabled: boolean; + metric?: QueryCostWarningMetric; + threshold?: number; +}; + export type DatabaseObject = { // Connection + general backend?: string; @@ -90,6 +98,7 @@ export type DatabaseObject = { // SQL Lab allows_cost_estimate?: boolean; cost_estimate_warning_threshold_bytes?: number | null; + query_cost_warning?: QueryCostWarning; allow_ctas?: boolean; allow_cvas?: boolean; allow_dml?: boolean; @@ -114,6 +123,7 @@ export type DatabaseObject = { // DB Engine Spec information engine_information?: { + query_cost_warning_metric?: QueryCostWarningMetric | null; supports_file_upload?: boolean; disable_ssh_tunneling?: boolean; supports_dynamic_catalog?: boolean; @@ -221,6 +231,7 @@ export type DatabaseForm = { preferred: boolean; sqlalchemy_uri_placeholder: string; engine_information: { + query_cost_warning_metric?: QueryCostWarningMetric | null; supports_file_upload: boolean; disable_ssh_tunneling: boolean; supports_dynamic_catalog: boolean; @@ -246,6 +257,7 @@ export interface ExtraJson { cost_estimate_enabled?: boolean; // in SQL Lab cost_estimate_warning_enabled?: boolean; cost_estimate_warning_threshold_bytes?: number; + query_cost_warning?: QueryCostWarning; disable_data_preview?: boolean; // in SQL Lab disable_drill_to_detail?: boolean; allow_multi_catalog?: boolean; From 64bad179f91db9664d00ed7cec8cc300dc67c220 Mon Sep 17 00:00:00 2001 From: Alankrit Verma Date: Sun, 18 Oct 2026 16:59:40 -0400 Subject: [PATCH 07/58] feat(sqllab): check advisory query cost warnings --- docs/docs/using-superset/exploring-data.mdx | 15 + superset/commands/sql_lab/estimate.py | 158 ++++- superset/commands/sql_lab/execute.py | 19 +- superset/sqllab/api.py | 7 +- superset/sqllab/schemas.py | 9 + superset/sqllab/sqllab_execution_context.py | 11 +- superset/sqllab/utils.py | 29 +- .../sql_lab/commands_tests.py | 25 + .../sql_lab/test_query_cost_warning.py | 620 ++++++++++++++++++ tests/unit_tests/sqllab/test_schemas.py | 45 ++ 10 files changed, 912 insertions(+), 26 deletions(-) create mode 100644 tests/unit_tests/commands/sql_lab/test_query_cost_warning.py diff --git a/docs/docs/using-superset/exploring-data.mdx b/docs/docs/using-superset/exploring-data.mdx index 7c0367d5154f..b96fe7e2e376 100644 --- a/docs/docs/using-superset/exploring-data.mdx +++ b/docs/docs/using-superset/exploring-data.mdx @@ -445,6 +445,21 @@ The existing manual estimate button keeps its formatted display. Warning checks are advisory and are intended for deliberate SQL Lab runs, not dashboards, scheduled queries, automatic retries, CTAS/CVAS, or enforced warehouse limits. +API clients can check a warning with `POST /api/v1/sqllab/estimate/`, setting +`check_warning` to `true` and passing the selected `sql`, `database_id`, `catalog`, +`schema`, `template_params`, and SQL Lab dropdown `query_limit`. The check applies +the execution row limits, including the extra row fetched to detect truncation. +The result reports `disabled`, `unsupported`, `within_threshold`, `exceeded`, or +`unavailable`, with the native metric, label, unit, raw value, threshold, and +message where applicable. Equality does not warn. Treat the server's status as +authoritative; byte counts can exceed JavaScript's safe integer range. + +Checking supports one read-only SELECT. Templates, multiple statements, and +custom SQL transformations are unavailable. A known estimation timeout or an +unusable estimate also returns unavailable. Invalid rules, access failures, and +database or network errors remain errors. Omitting `check_warning` preserves the +manual estimator response. + ### SQL Lab Tips **Schema and table browser**: The left-side table browser uses a collapsible treeview — click a schema to expand its tables, and click a table to see its columns and sample data inline. This makes navigating large schemas much faster than the previous flat list. diff --git a/superset/commands/sql_lab/estimate.py b/superset/commands/sql_lab/estimate.py index a6d86ea6b0d0..4bd9b2180f63 100644 --- a/superset/commands/sql_lab/estimate.py +++ b/superset/commands/sql_lab/estimate.py @@ -21,11 +21,14 @@ from flask import current_app as app from flask_babel import gettext as __ +from jinja2 import Environment from jinja2.exceptions import TemplateError +from marshmallow import ValidationError -from superset import is_feature_enabled, security_manager +from superset import config, db, is_feature_enabled, security_manager from superset.commands.base import BaseCommand from superset.daos.database import DatabaseDAO +from superset.databases.schemas import get_query_cost_warning from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import ( OAuth2RedirectError, @@ -39,7 +42,10 @@ ) from superset.jinja_context import get_template_processor from superset.models.core import Database +from superset.models.sql_lab import Query from superset.sql.parse import SQLScript +from superset.sql_lab import apply_limit +from superset.sqllab.utils import get_query_limit, set_query_limit from superset.utils import core as utils, json from superset.utils.rls import apply_rls @@ -53,6 +59,8 @@ class EstimateQueryCostType(TypedDict): catalog: str | None schema: str | None include_raw: NotRequired[bool] + check_warning: NotRequired[bool] + query_limit: NotRequired[int | None] class QueryEstimationCommand(BaseCommand): @@ -71,6 +79,8 @@ def __init__(self, params: EstimateQueryCostType) -> None: self._schema = params.get("schema") or "" self._catalog = params.get("catalog") self._include_raw = params.get("include_raw", False) + self._check_warning = params.get("check_warning", False) + self._query_limit = params.get("query_limit") def validate(self) -> None: # Load the database through the DAO so ``DatabaseFilter`` scopes @@ -167,11 +177,157 @@ def _apply_sql_security(self, sql: str) -> str: return sql + def _get_warning_rule(self, metric: str) -> dict[str, Any]: + """Validate saved settings without the display property's disabled fallback.""" + spec = self._database.db_engine_spec + try: + # Do not use Database.query_cost_warning: its display fallback hides + # invalid saved rules, which must remain errors on this path. + rule = get_query_cost_warning( + json.loads(self._database.extra or "{}"), spec.engine + ) + if rule["enabled"] and rule["metric"] != metric: + raise ValidationError( + { + "metric": [ + __("The warning metric is not supported by this engine.") + ] + } + ) + except ValidationError as ex: + raise SupersetErrorException( + SupersetError( + message=__("The saved query cost warning rule is invalid."), + error_type=SupersetErrorType.GENERIC_COMMAND_ERROR, + level=ErrorLevel.ERROR, + extra={"messages": ex.messages}, + ), + status=400, + ) from ex + return rule + + def _prepare_warning_sql(self) -> tuple[str | None, str | None]: + """Prepare execution-equivalent SQL, or explain why it cannot be estimated.""" + spec = self._database.db_engine_spec + if self._template_params or any( + kind != "data" for _, kind, _ in Environment(autoescape=True).lex(self._sql) + ): + return None, __("Warning checks are unavailable for SQL templates.") + mutator = app.config["SQL_QUERY_MUTATOR"] + if (mutator and mutator is not config.SQL_QUERY_MUTATOR) or ( + is_feature_enabled("ENABLE_TEMPLATE_PROCESSING") + and spec.engine in app.config["CUSTOM_TEMPLATE_PROCESSORS"] + ): + return None, __( + "Warning checks are unavailable for custom SQL transformations." + ) + + if self._catalog is None: + self._catalog = self._database.get_default_catalog() + sql = self._apply_sql_security(self._sql) + script = SQLScript(sql, engine=spec.engine) + if ( + len(script.statements) != 1 + or not script.statements[0].is_select() + or script.statements[0].is_mutating() + ): + return None, __("Warning checks require a single read-only SELECT.") + + # Use execution's dropdown/SQL limits and its extra truncation row. + # This Query is never persisted or dispatched for execution. + query = Query(database=self._database, select_as_cta_used=False) + if query in db.session: + db.session.expunge(query) + set_query_limit(query, sql, get_query_limit(self._query_limit)) + apply_limit(query, script.statements[0]) + return script.format(comments=spec.allows_sql_comments), None + + def _run_warning_check(self) -> dict[str, Any]: + """Compare raw usage while keeping failures distinct from advisory opt-outs.""" + if not ( + is_feature_enabled("SQLLAB_QUERY_COST_WARNINGS") + and is_feature_enabled("ESTIMATE_QUERY_COST") + ): + return {"status": "disabled"} + spec = self._database.db_engine_spec + metric = spec.get_query_cost_warning_metric() + if metric is None: + return {"status": "unsupported"} + if self._database.allows_cost_estimate is not True: + return {"status": "disabled"} + rule = self._get_warning_rule(metric) + if rule["enabled"] is False: + return {"status": "disabled"} + + label, unit = ( + (__("Estimated bytes processed"), __("bytes")) + if metric == "bytes_processed" + else (__("Estimated planner cost"), __("planner cost units")) + ) + result: dict[str, Any] = { + "status": "unavailable", + "metric": metric, + "label": label, + "unit": unit, + "value": None, + "threshold": rule["threshold"], + } + sql, reason = self._prepare_warning_sql() + if sql is None: + return {**result, "message": reason} + + timeout = app.config["SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT"] + try: + with utils.timeout( + seconds=timeout, + error_message=f"The estimation exceeded the {timeout} seconds timeout.", + ): + estimate = spec.estimate_query_cost_for_warning( + self._database, self._catalog, self._schema, sql + ) + except SupersetTimeoutException: + return {**result, "message": __("The query cost estimate timed out.")} + + comparison = estimate.compare(rule["threshold"]) + result.update( + accuracy=estimate.accuracy, + row_level_security_applied=estimate.row_level_security_applied, + ) + if estimate.row_level_security_applied is not True and estimate.accuracy in { + "PRECISE", + "LOWER_BOUND", + "UPPER_BOUND", + }: + result["value"] = estimate.value + if comparison == "unavailable": + return { + **result, + "message": __( + "The estimate is unavailable or cannot establish whether " + "usage is within the warning threshold." + ), + } + result["status"] = ( + "exceeded" if comparison == "above_threshold" else "within_threshold" + ) + result["message"] = ( + __("The estimate exceeds this connection's warning threshold.") + if result["status"] == "exceeded" + else __("The estimate is within this connection's warning threshold.") + ) + return result + def run( self, ) -> list[dict[str, Any]] | dict[str, Any]: self.validate() + if self._check_warning: + return self._run_warning_check() + return self._run_manual_estimate() + + def _run_manual_estimate(self) -> list[dict[str, Any]] | dict[str, Any]: + """Preserve the manual estimator's rendering, formatter and error contract.""" sql = self._sql if self._template_params: # Access is already checked in validate() before any rendering. diff --git a/superset/commands/sql_lab/execute.py b/superset/commands/sql_lab/execute.py index 99ed73fdef70..42ce134f5939 100644 --- a/superset/commands/sql_lab/execute.py +++ b/superset/commands/sql_lab/execute.py @@ -40,7 +40,7 @@ SqlLabException, ) from superset.sqllab.execution_context_convertor import ExecutionContextConvertor -from superset.sqllab.limiting_factor import LimitingFactor +from superset.sqllab.utils import set_query_limit from superset.utils.decorators import transaction if TYPE_CHECKING: @@ -251,21 +251,10 @@ def _is_required_to_set_limit(self) -> bool: ) def _set_query_limit(self, rendered_query: str) -> None: - db_engine_spec = self._execution_context.database.db_engine_spec # type: ignore - limits = [ - db_engine_spec.get_limit_from_sql(rendered_query), + set_query_limit( + self._execution_context.query, + rendered_query, self._execution_context.limit, - ] - if limits[0] is None or limits[0] > limits[1]: # type: ignore - self._execution_context.query.limiting_factor = LimitingFactor.DROPDOWN - elif limits[1] > limits[0]: # type: ignore - self._execution_context.query.limiting_factor = LimitingFactor.QUERY - else: # limits[0] == limits[1] - self._execution_context.query.limiting_factor = ( - LimitingFactor.QUERY_AND_DROPDOWN - ) - self._execution_context.query.limit = min( - lim for lim in limits if lim is not None ) diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 21a49feb149f..0425fdf9abfe 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -171,7 +171,12 @@ def estimate_query_cost(self) -> Response: $ref: '#/components/schemas/EstimateQueryCostSchema' responses: 200: - description: Query estimation result + description: >- + Formatted manual estimate, or an advisory status object when + check_warning is true. Warning statuses are disabled, unsupported, + within_threshold, exceeded, and unavailable. Active rules also return + metric, label, unit, value, threshold, and a short message. BigQuery + results retain accuracy and row_level_security_applied metadata. content: application/json: schema: diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py index 6e755e83a161..6e8ec47f5e32 100644 --- a/superset/sqllab/schemas.py +++ b/superset/sqllab/schemas.py @@ -51,6 +51,15 @@ class EstimateQueryCostSchema(Schema): allow_none=True, metadata={"description": "The database schema"} ) include_raw = fields.Boolean(load_default=False) + check_warning = fields.Boolean( + load_default=False, + metadata={"description": "Check the connection's advisory warning threshold"}, + ) + query_limit = fields.Integer( + allow_none=True, + strict=True, + metadata={"description": "The SQL Lab dropdown row limit"}, + ) class FormatQueryPayloadSchema(Schema): diff --git a/superset/sqllab/sqllab_execution_context.py b/superset/sqllab/sqllab_execution_context.py index 6e9bc7d00196..b9a9dfe84ab8 100644 --- a/superset/sqllab/sqllab_execution_context.py +++ b/superset/sqllab/sqllab_execution_context.py @@ -30,8 +30,9 @@ from superset.exceptions import SupersetErrorException from superset.models.sql_lab import Query from superset.sql.parse import CTASMethod +from superset.sqllab.utils import get_query_limit from superset.utils import core as utils, json -from superset.utils.core import apply_max_row_limit, get_user_id +from superset.utils.core import get_user_id from superset.utils.dates import now_as_float from superset.views.utils import get_cta_schema_name @@ -110,13 +111,7 @@ def _get_template_params(query_params: dict[str, Any]) -> dict[str, Any]: @staticmethod def _get_limit_param(query_params: dict[str, Any]) -> int: - limit = apply_max_row_limit(query_params.get("queryLimit") or 0) - if limit < 0: - logger.warning( - "Invalid limit of %i specified. Defaulting to max limit.", limit - ) - limit = 0 - return limit + return get_query_limit(query_params.get("queryLimit")) def is_run_asynchronous(self) -> bool: return self.async_flag diff --git a/superset/sqllab/utils.py b/superset/sqllab/utils.py index 81c7a3561f3a..e6c99cb7a222 100644 --- a/superset/sqllab/utils.py +++ b/superset/sqllab/utils.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import logging from typing import Any import pyarrow as pa @@ -23,7 +24,12 @@ from superset import db, is_feature_enabled from superset.common.db_query_status import QueryStatus from superset.daos.database import DatabaseDAO -from superset.models.sql_lab import TabState +from superset.models.sql_lab import Query, TabState +from superset.sqllab.limiting_factor import LimitingFactor +from superset.utils.core import apply_max_row_limit + +logger = logging.getLogger(__name__) + DATABASE_KEYS = [ "allow_file_upload", @@ -46,6 +52,27 @@ ] +def get_query_limit(query_limit: int | None) -> int: + """Normalize the SQL Lab dropdown limit using the execution row cap.""" + limit = apply_max_row_limit(query_limit or 0) + if limit < 0: + logger.warning("Invalid limit of %i specified. Defaulting to max limit.", limit) + limit = 0 + return limit + + +def set_query_limit(query: Query, rendered_query: str, limit: int) -> None: + """Choose the execution limit and record whether SQL or the dropdown sets it.""" + sql_limit = query.database.db_engine_spec.get_limit_from_sql(rendered_query) + if sql_limit is None or sql_limit > limit: + query.limiting_factor = LimitingFactor.DROPDOWN + elif limit > sql_limit: + query.limiting_factor = LimitingFactor.QUERY + else: + query.limiting_factor = LimitingFactor.QUERY_AND_DROPDOWN + query.limit = min(lim for lim in (sql_limit, limit) if lim is not None) + + def apply_display_max_row_configuration_if_require( # pylint: disable=invalid-name sql_results: dict[str, Any], max_rows_in_result: int ) -> dict[str, Any]: diff --git a/tests/integration_tests/sql_lab/commands_tests.py b/tests/integration_tests/sql_lab/commands_tests.py index 893d68d41aed..845e495446a2 100644 --- a/tests/integration_tests/sql_lab/commands_tests.py +++ b/tests/integration_tests/sql_lab/commands_tests.py @@ -49,6 +49,31 @@ class TestQueryEstimationCommand(SupersetTestCase): + @patch("superset.commands.sql_lab.estimate.is_feature_enabled", return_value=False) + def test_warning_preparation_does_not_pollute_session( + self, mock_is_feature_enabled: Mock + ) -> None: + """The limit probe must not leave an incomplete Query pending for autoflush.""" + database = get_example_database() + command = estimate.QueryEstimationCommand( + { + "database_id": database.id, + "sql": "SELECT 1", + "catalog": None, + "schema": None, + "template_params": {}, + "query_limit": 10, + "check_warning": True, + } + ) + command._database = database + sql, reason = command._prepare_warning_sql() + assert reason is None + assert sql is not None + assert "LIMIT 11" in sql + assert not any(isinstance(obj, Query) for obj in db.session.new) + db.session.flush() + def test_validation_no_database(self) -> None: params = {"database_id": 1, "sql": "SELECT 1"} schema = EstimateQueryCostSchema() diff --git a/tests/unit_tests/commands/sql_lab/test_query_cost_warning.py b/tests/unit_tests/commands/sql_lab/test_query_cost_warning.py new file mode 100644 index 000000000000..5e5fd13e0c1e --- /dev/null +++ b/tests/unit_tests/commands/sql_lab/test_query_cost_warning.py @@ -0,0 +1,620 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Warning checks must describe execution SQL without hiding errors or bounds.""" + +from typing import Any, cast +from unittest.mock import MagicMock + +import pytest +from flask import Flask +from pytest_mock import MockerFixture + +from superset import config +from superset.commands.sql_lab.estimate import ( + EstimateQueryCostType, + QueryEstimationCommand, +) +from superset.commands.sql_lab.execute import ExecuteSqlCommand +from superset.db_engine_specs.base import QueryCostEstimate +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import ( + OAuth2RedirectError, + SupersetDisallowedSQLFunctionException, + SupersetDisallowedSQLTableException, + SupersetDMLNotAllowedException, + SupersetErrorException, + SupersetParseError, + SupersetSecurityException, + SupersetTimeoutException, +) +from superset.models.core import Database +from superset.models.sql_lab import Query +from superset.sql.parse import SQLScript +from superset.sql_lab import apply_limit +from superset.sqllab.sqllab_execution_context import SqlJsonExecutionContext +from superset.utils import json + + +def _make_command( + mocker: MockerFixture, engine: str = "postgresql", **params: Any +) -> QueryEstimationCommand: + """Use real engine specs, SQL parsing, rule validation and execution limits.""" + metric = "bytes_processed" if engine == "bigquery" else "total_cost" + database = Database( + id=1, + database_name="warning_test", + sqlalchemy_uri=f"{engine}://localhost/test", + extra=json.dumps( + { + "cost_estimate_enabled": True, + "query_cost_warning": { + "enabled": True, + "metric": metric, + "threshold": 10, + }, + } + ), + allow_dml=False, + ) + parameters: dict[str, Any] = { + "database_id": 1, + "sql": "SELECT * FROM sales", + "catalog": "selected_catalog", + "schema": "selected_schema", + "template_params": {}, + "query_limit": 100, + "check_warning": True, + **params, + } + command = QueryEstimationCommand(cast(EstimateQueryCostType, parameters)) + mocker.patch( + "superset.commands.sql_lab.estimate.DatabaseDAO.find_by_id", + return_value=database, + ) + mocker.patch( + "superset.commands.sql_lab.estimate.security_manager", new_callable=MagicMock + ) + mocker.patch( + "superset.commands.sql_lab.estimate.is_feature_enabled", + side_effect=lambda name: name + in {"SQLLAB_QUERY_COST_WARNINGS", "ESTIMATE_QUERY_COST"}, + ) + mocker.patch.object(database, "get_default_catalog", return_value="default_catalog") + mocker.patch.object( + database.db_engine_spec, + "estimate_query_cost_for_warning", + return_value=QueryCostEstimate.from_value(5), + ) + return command + + +def _set_extra(command: QueryEstimationCommand, extra: dict[str, Any]) -> None: + """Replace raw saved settings so invalid data cannot use the display fallback.""" + command.validate() + command._database.extra = json.dumps(extra) + + +def _estimator(command: QueryEstimationCommand) -> MagicMock: + """Type the engine hook replaced by the test setup's synchronous mock.""" + return cast( + MagicMock, command._database.db_engine_spec.estimate_query_cost_for_warning + ) + + +def _run_warning(command: QueryEstimationCommand) -> dict[str, Any]: + """Narrow the endpoint's manual-or-warning result to the requested warning.""" + result = command.run() + assert isinstance(result, dict) + return result + + +@pytest.mark.parametrize("engine", ["bigquery", "postgresql"]) +@pytest.mark.parametrize( + "value,status", + [(0, "within_threshold"), (10, "within_threshold"), (11, "exceeded")], +) +def test_warning_comparison( + mocker: MockerFixture, engine: str, value: int, status: str +) -> None: + """Compare raw native values with strict greater-than, including valid zero.""" + command = _make_command(mocker, engine) + command.validate() + _estimator(command).return_value = QueryCostEstimate.from_value(value) + result = _run_warning(command) + assert isinstance(result, dict) + assert result["status"] == status + assert result["value"] == value + assert result["threshold"] == 10 + assert result["metric"] == ( + "bytes_processed" if engine == "bigquery" else "total_cost" + ) + assert result["unit"] == ("bytes" if engine == "bigquery" else "planner cost units") + assert result["label"] + assert result["message"] + + +def test_warning_compares_postgres_decimals_before_display( + mocker: MockerFixture, +) -> None: + """A planner estimate just above a decimal threshold must warn before rounding.""" + command = _make_command(mocker) + _set_extra( + command, + { + "cost_estimate_enabled": True, + "query_cost_warning": { + "enabled": True, + "metric": "total_cost", + "threshold": 10.125, + }, + }, + ) + _estimator(command).return_value = QueryCostEstimate.from_value(10.12501) + result = _run_warning(command) + assert result["status"] == "exceeded" + assert result["value"] == 10.12501 + + +@pytest.mark.parametrize( + "value,accuracy,masked,status", + [ + (10, "LOWER_BOUND", False, "unavailable"), + (9, "LOWER_BOUND", False, "unavailable"), + (11, "LOWER_BOUND", False, "exceeded"), + (10, "UPPER_BOUND", False, "within_threshold"), + (11, "UPPER_BOUND", False, "exceeded"), + (0, "UNKNOWN", False, "unavailable"), + (11, None, False, "unavailable"), + (11, "PRECISE", True, "unavailable"), + (2**63 - 1, "PRECISE", False, "exceeded"), + ], +) +def test_warning_preserves_bigquery_bounds( + mocker: MockerFixture, value: int, accuracy: str | None, masked: bool, status: str +) -> None: + """Unknown or masked bytes never imply safety; bounds keep their meaning.""" + command = _make_command(mocker, "bigquery") + command.validate() + _estimator(command).return_value = QueryCostEstimate.from_value( + value, accuracy, masked + ) + result = _run_warning(command) + assert result["status"] == status + assert result["accuracy"] == accuracy + assert result["row_level_security_applied"] is masked + assert result["value"] == ( + None if masked or accuracy in {None, "UNKNOWN"} else value + ) + + +def test_warning_compares_large_bytes_exactly(mocker: MockerFixture) -> None: + """A one-byte difference near JavaScript's integer boundary is compared here.""" + command = _make_command(mocker, "bigquery") + _set_extra( + command, + { + "cost_estimate_enabled": True, + "query_cost_warning": { + "enabled": True, + "metric": "bytes_processed", + "threshold": 2**53 - 1, + }, + }, + ) + _estimator(command).return_value = QueryCostEstimate.from_value(2**53) + assert _run_warning(command)["status"] == "exceeded" + + +@pytest.mark.parametrize("value", [None, True, "10", -1, float("nan"), float("inf")]) +def test_warning_unusable_estimate(mocker: MockerFixture, value: object) -> None: + """Missing and malformed estimates are unavailable, not within threshold.""" + command = _make_command(mocker) + command.validate() + _estimator(command).return_value = QueryCostEstimate.from_value(value) + assert _run_warning(command)["status"] == "unavailable" + + +@pytest.mark.parametrize("flag", ["SQLLAB_QUERY_COST_WARNINGS", "ESTIMATE_QUERY_COST"]) +def test_warning_flag_off_still_checks_access(mocker: MockerFixture, flag: str) -> None: + """Even disabled checks first run the existing database/table authorization.""" + command = _make_command(mocker) + flags = mocker.patch( + "superset.commands.sql_lab.estimate.is_feature_enabled", + side_effect=lambda name: name != flag, + ) + assert _run_warning(command) == {"status": "disabled"} + _estimator(command).assert_not_called() + manager = mocker.patch( + "superset.commands.sql_lab.estimate.security_manager", new_callable=MagicMock + ) + access = manager.raise_for_access + access.side_effect = SupersetSecurityException( + SupersetError( + message="Denied", + error_type=SupersetErrorType.TABLE_SECURITY_ACCESS_ERROR, + level=ErrorLevel.ERROR, + ) + ) + flags.reset_mock() + with pytest.raises(SupersetSecurityException): + _run_warning(command) + access.assert_called_once() + flags.assert_not_called() + + +@pytest.mark.parametrize( + "extra", + [ + { + "cost_estimate_enabled": False, + "query_cost_warning": { + "enabled": True, + "metric": "total_cost", + "threshold": 10, + }, + }, + { + "cost_estimate_enabled": True, + "query_cost_warning": { + "enabled": False, + "metric": "total_cost", + "threshold": 10, + }, + }, + {"cost_estimate_enabled": True}, + ], +) +def test_warning_connection_opt_out( + mocker: MockerFixture, extra: dict[str, Any] +) -> None: + """Connection estimation and explicit warning opt-in are independent switches.""" + command = _make_command(mocker) + _set_extra(command, extra) + assert _run_warning(command) == {"status": "disabled"} + _estimator(command).assert_not_called() + + +def test_warning_unsupported_engine(mocker: MockerFixture) -> None: + """Unsupported engines differ from supported engines with unestimatable SQL.""" + command = _make_command(mocker, "sqlite") + assert _run_warning(command) == {"status": "unsupported"} + _estimator(command).assert_not_called() + + +@pytest.mark.parametrize( + "rule", + [ + None, + {"enabled": "true"}, + {"enabled": True}, + {"enabled": True, "metric": "bytes_processed", "threshold": 10}, + {"enabled": True, "metric": "total_cost", "threshold": "10"}, + {"enabled": False, "threshold": -1}, + ], +) +def test_warning_invalid_rule_is_an_error(mocker: MockerFixture, rule: object) -> None: + """Never hide an invalid saved rule or fall back to the legacy draft.""" + command = _make_command(mocker) + _set_extra( + command, + { + "cost_estimate_enabled": True, + "query_cost_warning": rule, + "cost_estimate_warning_enabled": True, + "cost_estimate_warning_threshold_bytes": 1, + }, + ) + with pytest.raises(SupersetErrorException) as error: + _run_warning(command) + assert error.value.status == 400 + _estimator(command).assert_not_called() + + +@pytest.mark.parametrize( + "sql,params,reason", + [ + ("SELECT {{ amount }}", {}, "templates"), + ("SELECT 1", {"amount": 1}, "templates"), + ("SELECT 1 {# comment #}", {}, "templates"), + ("SELECT 1; SELECT 2", {}, "single read-only SELECT"), + ("SHOW TABLES", {}, "single read-only SELECT"), + ("CREATE TABLE t AS SELECT 1", {}, "single read-only SELECT"), + ("SELECT * INTO t FROM sales", {}, "single read-only SELECT"), + ( + "WITH deleted AS (DELETE FROM sales RETURNING *) SELECT * FROM deleted", + {}, + "single read-only SELECT", + ), + ], +) +def test_warning_unestimatable_sql( + mocker: MockerFixture, sql: str, params: dict[str, Any], reason: str +) -> None: + """Only a single read-only SELECT is eligible for the initial warning check.""" + command = _make_command(mocker, sql=sql, template_params=params) + command.validate() + command._database.allow_dml = True + result = _run_warning(command) + assert result["status"] == "unavailable" + assert reason in result["message"] + _estimator(command).assert_not_called() + + +@pytest.mark.parametrize("custom", ["SQL_QUERY_MUTATOR", "CUSTOM_TEMPLATE_PROCESSORS"]) +def test_warning_rejects_custom_transformations( + mocker: MockerFixture, app: Flask, custom: str +) -> None: + """Do not run arbitrary transformations to predict a future execution render.""" + command = _make_command(mocker) + transform = mocker.MagicMock() + mocker.patch.dict( + app.config, + { + custom: transform + if custom == "SQL_QUERY_MUTATOR" + else {"postgresql": transform} + }, + ) + mocker.patch( + "superset.commands.sql_lab.estimate.is_feature_enabled", + side_effect=lambda name: name != "RLS_IN_SQLLAB", + ) + assert _run_warning(command)["status"] == "unavailable" + transform.assert_not_called() + + +@pytest.mark.parametrize("engine", ["postgresql", "bigquery"]) +@pytest.mark.parametrize( + "sql_limit,dropdown,maximum,expected", + [ + (None, 100, 50000, 101), + (5, 100, 50000, 6), + (500, 100, 50000, 101), + (100, 100, 50000, 101), + (None, 100, 20, 21), + (None, None, 20, 21), + (None, 0, 20, 21), + (0, 100, 50000, 50001), + (None, -1, 20, 21), + ], +) +def test_warning_matches_execution_sql_limits( + mocker: MockerFixture, + app: Flask, + engine: str, + sql_limit: int | None, + dropdown: int | None, + maximum: int, + expected: int, +) -> None: + """Exercise the actual execution context, command and final extra-row rewrite.""" + sql = "SELECT * FROM sales" + ( # noqa: S608 + f" LIMIT {sql_limit}" if sql_limit is not None else "" + ) + command = _make_command(mocker, engine, sql=sql, query_limit=dropdown) + mocker.patch.dict( + app.config, + {"SQL_MAX_ROW": maximum, "SQL_QUERY_MUTATOR": config.SQL_QUERY_MUTATOR}, + ) + _run_warning(command) + database = command._database + estimator = _estimator(command) + prepared = estimator.call_args.args[3] + + context = SqlJsonExecutionContext( + {"database_id": 1, "sql": sql, "queryLimit": dropdown} + ) + context.set_database(database) + context.set_query(Query(database=database)) + execute = ExecuteSqlCommand( + context, + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + MagicMock(), + False, + ) + execute._set_query_limit_if_required(sql) + script = SQLScript(sql, engine=engine) + apply_limit(context.query, script.statements[0]) + assert prepared == script.format( + comments=database.db_engine_spec.allows_sql_comments + ) + assert ( + SQLScript(prepared, engine=engine).statements[0].get_limit_value() == expected + ) + assert estimator.call_args.args[1:3] == ("selected_catalog", "selected_schema") + + +def test_warning_resolves_execution_default_catalog(mocker: MockerFixture) -> None: + """Unspecified catalogs use the same database default as SQL Lab execution.""" + command = _make_command(mocker, catalog=None, schema=None) + _run_warning(command) + assert _estimator(command).call_args.args[1:3] == ("default_catalog", "") + + +def test_warning_estimates_rls_constrained_sql(mocker: MockerFixture) -> None: + """Security controls and RLS injection run before the final execution limit.""" + command = _make_command(mocker) + mocker.patch( + "superset.commands.sql_lab.estimate.is_feature_enabled", return_value=True + ) + command.validate() + mocker.patch.object( + command._database, + "resolve_query_default_schema", + return_value="selected_schema", + ) + + mocker.patch( + "superset.utils.rls.get_predicates_for_table", return_value=["tenant_id = 1"] + ) + _run_warning(command) + prepared = _estimator(command).call_args.args[3] + assert "tenant_id = 1" in prepared + assert "LIMIT 101" in prepared + + +@pytest.mark.parametrize( + "error", + [ + ConnectionError("network failed"), + TimeoutError("native network timeout"), + RuntimeError("unexpected failure"), + ValueError("native maximum bytes billed exceeded"), + OAuth2RedirectError("test", "tab", "https://example.invalid/oauth"), + ], +) +def test_warning_real_errors_propagate(mocker: MockerFixture, error: Exception) -> None: + """Native limits, network, OAuth and unexpected errors never become unavailable.""" + command = _make_command(mocker) + command.validate() + _estimator(command).side_effect = error + with pytest.raises(type(error)): + _run_warning(command) + + +def test_warning_known_timeout_is_unavailable(mocker: MockerFixture) -> None: + """Only the known estimation timeout is an advisory unavailable result.""" + command = _make_command(mocker) + command.validate() + _estimator(command).side_effect = SupersetTimeoutException( + SupersetErrorType.SQLLAB_TIMEOUT_ERROR, "timed out", ErrorLevel.ERROR + ) + result = _run_warning(command) + assert result["status"] == "unavailable" + assert "timed out" in result["message"] + + +@pytest.mark.parametrize( + "sql,setting,expected", + [ + ("DELETE FROM sales", {}, SupersetDMLNotAllowedException), + ( + "SELECT * FROM sales", + {"DISALLOWED_SQL_TABLES": {"postgresql": {"sales"}}}, + SupersetDisallowedSQLTableException, + ), + ( + "SELECT forbidden()", + {"DISALLOWED_SQL_FUNCTIONS": {"postgresql": {"forbidden"}}}, + SupersetDisallowedSQLFunctionException, + ), + ("SELECT * FROM", {}, SupersetParseError), + ], +) +def test_warning_sql_errors_are_not_unavailable( + mocker: MockerFixture, + app: Flask, + sql: str, + setting: dict[str, Any], + expected: type[Exception], +) -> None: + """Keep SQL parse errors and the existing execution controls as actual errors.""" + command = _make_command(mocker, sql=sql) + command.validate() + mocker.patch.object( + command._database, + "resolve_query_default_schema", + return_value="selected_schema", + ) + mocker.patch.dict(app.config, setting) + with pytest.raises(expected): + _run_warning(command) + _estimator(command).assert_not_called() + + +def test_warning_endpoint( + mocker: MockerFixture, client: Any, full_api_access: None +) -> None: + """The existing protected endpoint accepts the optional warning request fields.""" + _make_command(mocker) + response = client.post( + "/api/v1/sqllab/estimate/", + json={ + "database_id": 1, + "sql": "SELECT * FROM sales", + "catalog": "selected_catalog", + "schema": "selected_schema", + "template_params": {}, + "query_limit": 100, + "check_warning": True, + }, + ) + assert response.status_code == 200 + assert response.json["result"]["status"] == "within_threshold" + + +@pytest.mark.parametrize("limit", ["100", True, 1.5, {}]) +def test_warning_endpoint_rejects_invalid_input( + mocker: MockerFixture, client: Any, full_api_access: None, limit: object +) -> None: + """Malformed warning requests fail request validation before a warehouse call.""" + command = mocker.patch("superset.sqllab.api.QueryEstimationCommand") + response = client.post( + "/api/v1/sqllab/estimate/", + json={ + "database_id": 1, + "sql": "SELECT 1", + "check_warning": True, + "query_limit": limit, + }, + ) + assert response.status_code == 400 + command.assert_not_called() + + +def test_warning_endpoint_invalid_rule_is_error( + mocker: MockerFixture, client: Any, full_api_access: None +) -> None: + """Bad stored configuration is an HTTP error rather than a usable status.""" + command = _make_command(mocker) + _set_extra( + command, + {"cost_estimate_enabled": True, "query_cost_warning": {"enabled": True}}, + ) + response = client.post( + "/api/v1/sqllab/estimate/", + json={"database_id": 1, "sql": "SELECT 1", "check_warning": True}, + ) + assert response.status_code == 400 + assert "status" not in response.json + + +def test_warning_endpoint_keeps_manual_response( + mocker: MockerFixture, client: Any, full_api_access: None +) -> None: + """Requests without check_warning retain the manual list and formatter inputs.""" + command = _make_command(mocker) + command.validate() + spec = command._database.db_engine_spec + estimate = mocker.patch.object( + spec, "estimate_query_cost", return_value=[{"Total cost": 5.125}] + ) + formatter = mocker.patch.object( + spec, "query_cost_formatter", return_value=[{"Total cost": "5.1"}] + ) + response = client.post( + "/api/v1/sqllab/estimate/", json={"database_id": 1, "sql": "SELECT 1"} + ) + assert response.status_code == 200 + assert response.json == {"result": [{"Total cost": "5.1"}]} + formatter.assert_called_once_with([{"Total cost": 5.125}]) + estimate.assert_called_once() + _estimator(command).assert_not_called() diff --git a/tests/unit_tests/sqllab/test_schemas.py b/tests/unit_tests/sqllab/test_schemas.py index 2c924c077a38..d0257a01a794 100644 --- a/tests/unit_tests/sqllab/test_schemas.py +++ b/tests/unit_tests/sqllab/test_schemas.py @@ -14,6 +14,51 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import pytest +from marshmallow import ValidationError + +from superset.sqllab.schemas import EstimateQueryCostSchema + + +def test_estimate_schema_keeps_manual_defaults() -> None: + """Existing clients need neither warning checks nor a dropdown limit.""" + assert EstimateQueryCostSchema().load({"database_id": 1, "sql": "SELECT 1"}) == { + "database_id": 1, + "sql": "SELECT 1", + "include_raw": False, + "check_warning": False, + } + + +@pytest.mark.parametrize("limit", [None, 0, 100, -1]) +def test_estimate_schema_accepts_execution_limits(limit: int | None) -> None: + """Integer limits use the same normalization as SQL Lab execution.""" + result = EstimateQueryCostSchema().load( + { + "database_id": 1, + "sql": "SELECT 1", + "check_warning": True, + "query_limit": limit, + } + ) + assert result["query_limit"] == limit + assert result["check_warning"] is True + + +@pytest.mark.parametrize("limit", ["100", True, 1.5, {}, float("inf")]) +def test_estimate_schema_rejects_invalid_warning_limit(limit: object) -> None: + """Do not silently coerce malformed request limits before checking usage.""" + with pytest.raises(ValidationError, match="Not a valid integer"): + EstimateQueryCostSchema().load( + { + "database_id": 1, + "sql": "SELECT 1", + "check_warning": True, + "query_limit": limit, + } + ) + + def test_sqllab_bootstrap_database_schema_includes_engine_information() -> None: """ The OpenAPI contract for `GET /api/v1/sqllab/` documents `engine_information` From 606c941669d0b406c085263b9bae0bbf56661195 Mon Sep 17 00:00:00 2001 From: Alankrit Verma Date: Sun, 18 Oct 2026 17:38:54 -0400 Subject: [PATCH 08/58] feat(database): restore warning units with exact byte validation --- docs/docs/using-superset/exploring-data.mdx | 7 + .../DatabaseModal.queryCostWarning.test.tsx | 203 ++++++++++++ .../DatabaseModal/ExtraOptions.test.tsx | 19 +- .../databases/DatabaseModal/ExtraOptions.tsx | 84 +---- .../QueryCostWarningSettings.test.tsx | 306 ++++++++++++++++++ .../QueryCostWarningSettings.tsx | 240 ++++++++++++++ .../databases/DatabaseModal/index.test.tsx | 2 +- .../databases/DatabaseModal/index.tsx | 37 ++- .../DatabaseModal/queryCostWarning.test.ts | 79 +++++ .../DatabaseModal/queryCostWarning.ts | 93 +++++- .../databases/DatabaseModal/styles.ts | 10 +- 11 files changed, 989 insertions(+), 91 deletions(-) create mode 100644 superset-frontend/src/features/databases/DatabaseModal/DatabaseModal.queryCostWarning.test.tsx create mode 100644 superset-frontend/src/features/databases/DatabaseModal/QueryCostWarningSettings.test.tsx create mode 100644 superset-frontend/src/features/databases/DatabaseModal/QueryCostWarningSettings.tsx diff --git a/docs/docs/using-superset/exploring-data.mdx b/docs/docs/using-superset/exploring-data.mdx index b96fe7e2e376..920a6a210ebb 100644 --- a/docs/docs/using-superset/exploring-data.mdx +++ b/docs/docs/using-superset/exploring-data.mdx @@ -424,6 +424,13 @@ positive number that can include decimals. PostgreSQL planner units are not runtime, money, or memory. BigQuery estimates describe usage, not a guaranteed bill. These metrics are not converted into a common unit. +BigQuery thresholds can be entered in Bytes, MB, or GB (decimal units: 1 MB = +1,000,000 bytes and 1 GB = 1,000,000,000 bytes). Changing the unit changes the +display while retaining the saved byte amount. Fractional MB or GB are accepted +only when they represent a positive whole number of bytes in the allowed range; +invalid amounts show validation instead of being rounded. PostgreSQL thresholds +stay in relative planner cost units and have no default amount. + Disabling warnings retains valid saved values. Turning connection estimation off also retains the rule and makes checking inactive. A broken rule can be replaced with `{ "enabled": false }` to disable it without retaining invalid fields. diff --git a/superset-frontend/src/features/databases/DatabaseModal/DatabaseModal.queryCostWarning.test.tsx b/superset-frontend/src/features/databases/DatabaseModal/DatabaseModal.queryCostWarning.test.tsx new file mode 100644 index 000000000000..a33cb4abcc70 --- /dev/null +++ b/superset-frontend/src/features/databases/DatabaseModal/DatabaseModal.queryCostWarning.test.tsx @@ -0,0 +1,203 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import fetchMock from 'fetch-mock'; +import { FeatureFlag } from '@superset-ui/core'; +import { + fireEvent, + render, + screen, + userEvent, + waitFor, +} from 'spec/helpers/testing-library'; +import DatabaseModal from './index'; +import { DatabaseObject, ExtraJson, QueryCostWarningMetric } from '../types'; + +const databaseEndpoint = 'glob:*/api/v1/database/10'; +const connectionEndpoint = 'glob:*/api/v1/database/10/connection'; + +beforeEach(() => { + window.featureFlags = { + [FeatureFlag.EstimateQueryCost]: true, + [FeatureFlag.SqlLabQueryCostWarnings]: true, + }; + fetchMock.get('glob:*/api/v1/database/available*', { databases: [] }); +}); +afterEach(() => fetchMock.clearHistory().removeRoutes()); + +test.each([ + { + engine: 'postgresql', + metric: 'total_cost' as QueryCostWarningMetric, + threshold: 12.5, + input: '14.125', + saved: 14.125, + reopened: 14.125, + label: 'Warning threshold (planner cost)', + }, + { + engine: 'bigquery', + metric: 'bytes_processed' as QueryCostWarningMetric, + threshold: 2000000000, + input: '1.25', + saved: 1250000000, + reopened: 1250, + label: 'Warning threshold (bytes processed)', + }, +])( + 'saves and reopens disabled settings for $engine', + async ({ engine, metric, threshold, input, saved, reopened, label }) => { + const initialExtra = { + cost_estimate_enabled: true, + allows_virtual_table_explore: true, + query_cost_warning: { enabled: true, metric, threshold }, + }; + let database = { + id: 10, + database_name: 'Warning connection', + backend: engine, + expose_in_sqllab: true, + configuration_method: 'sqlalchemy_form', + sqlalchemy_uri: `${engine}://localhost/test`, + engine_information: { query_cost_warning_metric: metric }, + extra: JSON.stringify(initialExtra), + }; + fetchMock.get(connectionEndpoint, () => ({ result: database })); + fetchMock.put( + databaseEndpoint, + { id: 10, result: database }, + { name: 'database-update' }, + ); + const view = render( + , + { useRedux: true }, + ); + await userEvent.click(await screen.findByRole('tab', { name: 'Advanced' })); + await userEvent.click(await screen.findByText('SQL Lab')); + const amount = await screen.findByLabelText(label); + fireEvent.change(amount, { target: { value: input } }); + await userEvent.click( + screen.getByLabelText('Enable pre-run cost warnings'), + ); + await userEvent.click(screen.getByRole('button', { name: 'Finish' })); + await waitFor(() => + expect(fetchMock.callHistory.calls('database-update')).toHaveLength(1), + ); + const request = JSON.parse( + fetchMock.callHistory.calls('database-update')[0].options.body as string, + ) as DatabaseObject; + const extra = JSON.parse(request.extra || '{}') as ExtraJson; + expect(extra).toMatchObject({ + ...initialExtra, + query_cost_warning: { enabled: false, metric, threshold: saved }, + }); + view.unmount(); + database = { ...database, extra: request.extra || '{}' }; + render(, { + useRedux: true, + }); + await userEvent.click(await screen.findByRole('tab', { name: 'Advanced' })); + await userEvent.click(await screen.findByText('SQL Lab')); + expect(await screen.findByLabelText(label)).toHaveValue(reopened); + expect( + screen.getByLabelText('Enable pre-run cost warnings'), + ).not.toBeChecked(); + }, +); + +test('the modal blocks an invalid byte amount and allows correcting it', async () => { + fetchMock.get(connectionEndpoint, { + result: { + id: 10, + database_name: 'Byte validation', + backend: 'bigquery', + expose_in_sqllab: true, + configuration_method: 'sqlalchemy_form', + engine_information: { query_cost_warning_metric: 'bytes_processed' }, + extra: JSON.stringify({ + cost_estimate_enabled: true, + query_cost_warning: { + enabled: true, + metric: 'bytes_processed', + threshold: 1000000, + }, + }), + }, + }); + render(, { + useRedux: true, + }); + await userEvent.click(await screen.findByRole('tab', { name: 'Advanced' })); + await userEvent.click(await screen.findByText('SQL Lab')); + const amount = await screen.findByLabelText( + 'Warning threshold (bytes processed)', + ); + fireEvent.change(amount, { target: { value: '0.0000001' } }); + expect(screen.getByRole('button', { name: 'Finish' })).toBeDisabled(); + expect( + screen.getByText(/amount that equals positive whole bytes/), + ).toBeInTheDocument(); + fireEvent.change(amount, { target: { value: '1.000001' } }); + expect(screen.getByRole('button', { name: 'Finish' })).toBeEnabled(); +}); + +test('saving legacy bytes writes the new disabled rule without implicit opt-in', async () => { + const extra = { + cost_estimate_warning_threshold_bytes: 1500000, + cost_estimate_enabled: true, + allows_virtual_table_explore: true, + }; + const database = { + id: 10, + database_name: 'Legacy bytes', + backend: 'bigquery', + expose_in_sqllab: true, + configuration_method: 'sqlalchemy_form', + engine_information: { query_cost_warning_metric: 'bytes_processed' }, + extra: JSON.stringify(extra), + }; + fetchMock.get(connectionEndpoint, { result: database }); + fetchMock.put( + databaseEndpoint, + { id: 10, result: database }, + { name: 'legacy-update' }, + ); + render(, { + useRedux: true, + }); + await userEvent.click(await screen.findByRole('tab', { name: 'Advanced' })); + await userEvent.click(await screen.findByText('SQL Lab')); + expect( + await screen.findByLabelText('Warning threshold (bytes processed)'), + ).toHaveValue(1500000); + await userEvent.click(screen.getByRole('button', { name: 'Finish' })); + await waitFor(() => + expect(fetchMock.callHistory.calls('legacy-update')).toHaveLength(1), + ); + const request = JSON.parse( + fetchMock.callHistory.calls('legacy-update')[0].options.body as string, + ) as DatabaseObject; + expect(JSON.parse(request.extra || '{}')).toMatchObject({ + ...extra, + query_cost_warning: { + enabled: false, + metric: 'bytes_processed', + threshold: 1500000, + }, + }); +}); diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx index a4ae5e668216..207cc72a43b8 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.test.tsx @@ -170,16 +170,20 @@ describe('ExtraOptions Component', () => { metric: 'bytes_processed' as const, label: 'Warning threshold (bytes processed)', threshold: 100000000, + display: 100, + updated: '123500000', }, { backend: 'postgresql', metric: 'total_cost' as const, label: 'Warning threshold (planner cost)', threshold: 12.5, + display: 12.5, + updated: '123.5', }, ])( 'warning thresholds retain native units for $backend', - ({ backend, metric, label, threshold }) => { + ({ backend, metric, label, threshold, display, updated }) => { renderComponent({ ...defaultDb, backend, @@ -194,13 +198,13 @@ describe('ExtraOptions Component', () => { }); fireEvent.click(screen.getByText('SQL Lab')); const amount = screen.getByLabelText(label); - expect(amount).toHaveValue(threshold); + expect(amount).toHaveValue(display); fireEvent.change(amount, { target: { value: '123.5' } }); expect(onExtraInputChange).toHaveBeenCalledWith( expect.objectContaining({ target: expect.objectContaining({ name: 'query_cost_warning_threshold', - value: '123.5', + value: updated, }), }), ); @@ -256,7 +260,7 @@ describe('ExtraOptions Component', () => { expect(onExtraInputChange).toHaveBeenCalled(); }); - test('warning settings are hidden when the warning flag is off', () => { + test('warning settings explain when the warning flag is off', () => { window.featureFlags = { [FeatureFlag.EstimateQueryCost]: true }; renderComponent({ ...defaultDb, @@ -267,8 +271,11 @@ describe('ExtraOptions Component', () => { }); fireEvent.click(screen.getByText('SQL Lab')); expect( - screen.queryByLabelText('Enable pre-run cost warnings'), - ).not.toBeInTheDocument(); + screen.getByLabelText('Enable pre-run cost warnings'), + ).toBeDisabled(); + expect( + screen.getByText(/SQLLAB_QUERY_COST_WARNINGS is off/), + ).toBeInTheDocument(); }); test('warnings cannot be enabled when the global estimate flag is off', () => { diff --git a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx index 7135662bedfa..a5ce3273e47c 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx @@ -31,7 +31,6 @@ import { Collapse, InfoTooltip, CollapseLabelInModal, - Flex, type CheckboxChangeEvent, } from '@superset-ui/core/components'; import { useJsonValidation } from '@superset-ui/core/components/AsyncAceEditor'; @@ -42,7 +41,7 @@ import { no_margin_bottom, } from './styles'; import { DatabaseObject, ExtraJson } from '../types'; -import { getQueryCostWarning } from './queryCostWarning'; +import QueryCostWarningSettings from './QueryCostWarningSettings'; const ExtraOptions = ({ db, @@ -52,6 +51,7 @@ const ExtraOptions = ({ onExtraInputChange, onExtraEditorChange, extraExtension, + onQueryCostWarningValidationChange, }: { db: DatabaseObject | null; onInputChange: ( @@ -64,6 +64,7 @@ const ExtraOptions = ({ ) => void; onExtraEditorChange: Function; extraExtension: DatabaseConnectionExtension | undefined; + onQueryCostWarningValidationChange?: (hasError: boolean) => void; }) => { const onExtraInputChangeNonNegative = (e: ChangeEvent) => { const { name, value } = e.target; @@ -103,19 +104,6 @@ const ExtraOptions = ({ return value; }); - const warningMetric = db?.engine_information?.query_cost_warning_metric; - const warningRule = getQueryCostWarning(extraJson, db?.backend || db?.engine); - const warningsEnabled = warningRule.enabled === true; - const warningAmount = - warningRule.metric === warningMetric && - typeof warningRule.threshold === 'number' && - Number.isFinite(warningRule.threshold) - ? warningRule.threshold - : ''; - const estimationEnabled = - isFeatureEnabled(FeatureFlag.EstimateQueryCost) && - extraJson.cost_estimate_enabled === true; - // JSON validation hooks for the three editors const secureExtraAnnotations = useJsonValidation(db?.masked_encrypted_extra, { errorPrefix: 'Invalid secure extra JSON', @@ -207,7 +195,6 @@ const ExtraOptions = ({ @@ -303,63 +290,14 @@ const ExtraOptions = ({ />
- {isFeatureEnabled(FeatureFlag.SqlLabQueryCostWarnings) && - (warningMetric || 'query_cost_warning' in extraJson) && ( - -
- - {t('Enable pre-run cost warnings')} - -
- {warningMetric && - (warningsEnabled || warningAmount !== '') && ( - - - - - {warningMetric === 'bytes_processed' - ? t( - 'Estimated bytes processed are usage estimates, not a guaranteed bill. You can still choose to run the query.', - ) - : t( - 'PostgreSQL planner cost uses planner units, not runtime or money. You can still choose to run the query.', - )} - - - )} - - )} +
+ { + if (typeof target.name !== 'string') + throw new Error('Expected a named setting'); + dispatch({ + type: ActionType.ExtraInputChange, + payload: { + name: target.name, + type: target.type, + value: target.value, + checked: 'checked' in target ? target.checked : false, + }, + }); + }} + /> + + + {db?.extra} + + ); +} + +/** Read only the serialized settings that the database modal would submit. */ +function savedExtra(): ExtraJson { + return JSON.parse( + screen.getByTestId('saved-extra').textContent || '{}', + ) as ExtraJson; +} + +beforeEach(() => { + window.featureFlags = { + [FeatureFlag.EstimateQueryCost]: true, + [FeatureFlag.SqlLabQueryCostWarnings]: true, + }; +}); + +test('unit changes preserve bytes and disabled values survive reopening', async () => { + const initial: ExtraJson = { + cost_estimate_enabled: true, + cost_estimate_warning_threshold_bytes: 2000000000, + cost_estimate_warning_enabled: true, + allows_virtual_table_explore: true, + query_cost_warning: { + enabled: true, + metric: 'bytes_processed', + threshold: 2000000000, + }, + }; + const view = render(); + const amount = screen.getByLabelText('Warning threshold (bytes processed)'); + expect(amount).toHaveValue(2); + await selectOption('MB', 'Threshold size unit'); + expect(amount).toHaveValue(2000); + expect(savedExtra()).toEqual(initial); + await selectOption('Bytes', 'Threshold size unit'); + expect(amount).toHaveValue(2000000000); + expect(savedExtra()).toEqual(initial); + await selectOption('GB', 'Threshold size unit'); + fireEvent.change(amount, { target: { value: '1.25' } }); + expect(savedExtra().query_cost_warning?.threshold).toBe(1250000000); + await userEvent.click(screen.getByLabelText('Enable pre-run cost warnings')); + expect(amount).toHaveValue(1.25); + expect(amount).toBeDisabled(); + const saved = savedExtra(); + expect(saved).toEqual({ + ...initial, + query_cost_warning: { + enabled: false, + metric: 'bytes_processed', + threshold: 1250000000, + }, + }); + view.unmount(); + render(); + expect( + screen.getByLabelText('Enable pre-run cost warnings'), + ).not.toBeChecked(); + expect( + screen.getByLabelText('Warning threshold (bytes processed)'), + ).toHaveValue(1250); + expect(savedExtra()).toEqual(saved); +}); + +test('fractional bytes show validation and cannot be saved or silently rounded', async () => { + render( + , + ); + const amount = screen.getByLabelText('Warning threshold (bytes processed)'); + fireEvent.change(amount, { target: { value: '0.0000001' } }); + expect(amount).toHaveValue(0.0000001); + expect( + screen.getByText(/amount that equals positive whole bytes/), + ).toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled(); + expect(savedExtra().query_cost_warning?.threshold).toBeUndefined(); + fireEvent.change(amount, { target: { value: '1.000001' } }); + expect(savedExtra().query_cost_warning?.threshold).toBe(1000001); + expect(screen.getByRole('button', { name: 'Save' })).toBeEnabled(); +}); + +test('large byte amounts display exactly in every unit', async () => { + render( + , + ); + const amount = screen.getByLabelText('Warning threshold (bytes processed)'); + await selectOption('GB', 'Threshold size unit'); + expect(amount).toHaveAttribute('value', '9007199.254740991'); + expect(savedExtra().query_cost_warning?.threshold).toBe( + Number.MAX_SAFE_INTEGER, + ); + fireEvent.change(amount, { target: { value: '9007199.254740992' } }); + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled(); + fireEvent.change(amount, { target: { value: '9007199.254740991' } }); + expect(savedExtra().query_cost_warning?.threshold).toBe( + Number.MAX_SAFE_INTEGER, + ); +}); + +test('turning connection estimation off retains the complete warning rule', async () => { + const rule = { + enabled: true, + metric: 'bytes_processed' as const, + threshold: 1500000, + }; + render( + , + ); + await userEvent.click( + screen.getByRole('button', { name: 'Turn estimation off' }), + ); + expect(screen.getByText(/Saved settings are retained/)).toBeInTheDocument(); + expect(savedExtra().query_cost_warning).toEqual(rule); + expect(screen.getByLabelText('Enable pre-run cost warnings')).toBeChecked(); + expect( + screen.getByLabelText('Warning threshold (bytes processed)'), + ).toBeDisabled(); + expect(screen.getByRole('button', { name: 'Save' })).toBeEnabled(); +}); + +test('PostgreSQL accepts decimals and explains its relative planner cost', () => { + render( + , + ); + const amount = screen.getByLabelText('Warning threshold (planner cost)'); + expect(amount).toHaveValue(null); + expect( + screen.getByText(/relative, not dollars or memory/), + ).toBeInTheDocument(); + expect(screen.queryByRole('combobox')).not.toBeInTheDocument(); + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled(); + fireEvent.change(amount, { target: { value: '12.125' } }); + expect(savedExtra().query_cost_warning?.threshold).toBe(12.125); + expect(screen.getByRole('button', { name: 'Save' })).toBeEnabled(); +}); + +test.each([undefined, false, true])( + 'legacy values prefill with explicit opt-in %s', + enabled => { + const extra = { + cost_estimate_enabled: true, + cost_estimate_warning_enabled: enabled, + cost_estimate_warning_threshold_bytes: 1500000, + }; + render(); + expect( + screen.getByLabelText('Warning threshold (bytes processed)'), + ).toHaveValue(1500000); + expect( + screen.getByLabelText('Enable pre-run cost warnings'), + ).toHaveProperty('checked', enabled === true); + expect(savedExtra()).toEqual(extra); + }, +); + +test.each([ + null, + { enabled: 'invalid', metric: 'bytes_processed', threshold: 10 }, + { enabled: false, metric: 'bytes_processed', threshold: -1 }, +])('broken saved rules remain disableable (%j)', rule => { + window.featureFlags = {}; + render(); + expect(screen.getByRole('button', { name: 'Save' })).toBeDisabled(); + fireEvent.click(screen.getByRole('button', { name: 'Disable saved rule' })); + expect(savedExtra().query_cost_warning?.enabled).toBe(false); + expect(screen.getByRole('button', { name: 'Save' })).toBeEnabled(); +}); + +test('a no-longer-supported saved rule can be disabled while preserving its value', () => { + render( + , + ); + expect(screen.getByText(/This engine does not support/)).toBeInTheDocument(); + fireEvent.click(screen.getByLabelText('Enable pre-run cost warnings')); + expect(savedExtra().query_cost_warning).toEqual({ + enabled: false, + metric: 'bytes_processed', + threshold: 100, + }); + expect(screen.getByRole('button', { name: 'Save' })).toBeEnabled(); +}); diff --git a/superset-frontend/src/features/databases/DatabaseModal/QueryCostWarningSettings.tsx b/superset-frontend/src/features/databases/DatabaseModal/QueryCostWarningSettings.tsx new file mode 100644 index 000000000000..3e531582f300 --- /dev/null +++ b/superset-frontend/src/features/databases/DatabaseModal/QueryCostWarningSettings.tsx @@ -0,0 +1,240 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { ChangeEvent, useEffect, useState } from 'react'; +import { t } from '@apache-superset/core/translation'; +import { useTheme } from '@apache-superset/core/theme'; +import { FeatureFlag, isFeatureEnabled } from '@superset-ui/core'; +import { + Button, + Checkbox, + CheckboxChangeEvent, + Flex, + Form, + Input, + Select, +} from '@superset-ui/core/components'; +import { ExtraJson, QueryCostWarningMetric } from '../types'; +import { + CostWarningUnit, + costWarningUnits, + formatWarningBytes, + getCostWarningUnit, + getQueryCostWarning, + getQueryCostWarningError, + parseWarningBytes, +} from './queryCostWarning'; +import { no_margin_bottom, StyledInputContainer } from './styles'; + +/** Edit a native warning threshold while keeping display units out of saved extras. */ +export default function QueryCostWarningSettings({ + extra, + engine, + metric, + onChange, + onValidationChange, +}: { + extra: ExtraJson; + engine?: string; + metric?: QueryCostWarningMetric | null; + onChange: ( + event: CheckboxChangeEvent | ChangeEvent, + ) => void; + onValidationChange?: (hasError: boolean) => void; +}) { + const theme = useTheme(); + const rule = getQueryCostWarning(extra, engine); + const enabled = rule.enabled === true; + const flagsEnabled = isFeatureEnabled(FeatureFlag.SqlLabQueryCostWarnings); + const estimationFlag = isFeatureEnabled(FeatureFlag.EstimateQueryCost); + const active = + flagsEnabled && + estimationFlag && + extra.cost_estimate_enabled === true && + !!metric; + const [unit, setUnit] = useState(() => + getCostWarningUnit(rule.threshold), + ); + const [draft, setDraft] = useState(); + const matchingThreshold = + rule.metric === metric && + typeof rule.threshold === 'number' && + Number.isFinite(rule.threshold) + ? rule.threshold + : undefined; + const value = + draft ?? + (matchingThreshold === undefined + ? '' + : metric === 'bytes_processed' + ? formatWarningBytes(matchingThreshold, unit) + : String(matchingThreshold)); + const draftInvalid = + draft !== undefined && + (metric === 'bytes_processed' + ? parseWarningBytes(draft, unit) === undefined + : draft.trim() === '' || + !Number.isFinite(Number(draft)) || + Number(draft) <= 0); + const error = draftInvalid + ? metric === 'bytes_processed' + ? t( + 'Enter an amount that equals positive whole bytes, no greater than 9,007,199,254,740,991.', + ) + : t('Enter a positive, finite warning threshold.') + : getQueryCostWarningError(rule, metric); + useEffect(() => { + onValidationChange?.(!!error); + }, [error, onValidationChange]); + + const setEnabled = (event: CheckboxChangeEvent) => { + setDraft(undefined); + onChange(event); + }; + const changeAmount = (event: ChangeEvent) => { + const input = event.target.value; + setDraft(input); + const threshold = + metric === 'bytes_processed' + ? parseWarningBytes(input, unit) + : Number(input); + onChange({ + target: { + name: 'query_cost_warning_threshold', + type: 'number', + value: + threshold === undefined || + !Number.isFinite(threshold) || + threshold <= 0 + ? '' + : String(threshold), + }, + } as ChangeEvent); + }; + const inactiveReason = !metric + ? t( + 'This engine does not support pre-run warnings. BigQuery and PostgreSQL are supported.', + ) + : !flagsEnabled + ? t('Warnings are inactive because SQLLAB_QUERY_COST_WARNINGS is off.') + : !estimationFlag + ? t('Warnings are inactive because ESTIMATE_QUERY_COST is off.') + : extra.cost_estimate_enabled !== true + ? t( + 'Enable query cost estimation to activate pre-run warnings. Saved settings are retained.', + ) + : undefined; + + return ( + +
+ + {t('Enable pre-run cost warnings')} + +
+ + {inactiveReason && {inactiveReason}} + {metric && (enabled || matchingThreshold !== undefined) && ( + + + + + {metric === 'bytes_processed' && ( + - {costWarningContext} + ); }; diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.test.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.test.tsx new file mode 100644 index 000000000000..7955faa595e4 --- /dev/null +++ b/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.test.tsx @@ -0,0 +1,493 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { applyMiddleware, createStore } from 'redux'; +import thunk from 'redux-thunk'; +import { renderHook, act } from '@testing-library/react'; +import { FeatureFlag, isFeatureEnabled } from '@superset-ui/core'; +import fetchMock from 'fetch-mock'; +import { createWrapper, waitFor } from 'spec/helpers/testing-library'; +import { initialState, defaultQueryEditor } from 'src/SqlLab/fixtures'; +import { DEFAULT_BOOTSTRAP_DATA } from 'src/constants'; +import type { SqlLabRootState } from 'src/SqlLab/types'; +import useQueryCostWarning, { parseWarningResult } from './useQueryCostWarning'; + +jest.mock('@superset-ui/core', () => ({ + ...jest.requireActual('@superset-ui/core'), + isFeatureEnabled: jest.fn(), +})); +const featureEnabled = jest.mocked(isFeatureEnabled); +const endpoint = 'glob:*/api/v1/sqllab/estimate/'; +const executeEndpoint = 'glob:*/api/v1/sqllab/execute/*'; +const warning = { + status: 'exceeded', + metric: 'total_cost', + label: 'Estimated planner cost', + unit: 'planner cost units', + value: 12.5, + threshold: 10, + message: 'Above threshold', +}; + +/** Control an estimate independently of the editor and execution response. */ +function setup() { + const resolvers: Array<(result: object) => void> = []; + const responses = [0, 1].map( + index => + new Promise(resolve => { + resolvers[index] = resolve; + }), + ); + let requestIndex = 0; + fetchMock.post( + endpoint, + () => { + const response = responses[requestIndex]; + requestIndex += 1; + return response; + }, + { + name: 'warning-check', + }, + ); + fetchMock.post( + executeEndpoint, + { data: [], columns: [] }, + { name: 'warning-execute' }, + ); + const editor = { + ...defaultQueryEditor, + sql: 'SELECT 1', + selectedText: '', + templateParams: '{}', + queryLimit: 0, + }; + const state: SqlLabRootState = { + ...initialState, + localStorageUsageInKilobytes: 0, + user: { + ...initialState.user, + isAnonymous: false, + groups: [], + roles: { Admin: [] }, + }, + common: { + ...DEFAULT_BOOTSTRAP_DATA.common, + conf: initialState.common.conf, + }, + sqlLab: { + ...initialState.sqlLab, + dbConnect: false, + errorMessage: null, + editorTabLastUpdatedAt: 0, + lastUpdatedActiveTab: '', + queryEditors: [editor], + unsavedQueryEditor: {}, + tabHistory: [editor.id], + databases: { + [editor.dbId!]: { + id: editor.dbId, + allow_run_async: false, + allows_cost_estimate: true, + query_cost_warning: { + enabled: true, + metric: 'total_cost', + threshold: 10, + }, + engine_information: { query_cost_warning_metric: 'total_cost' }, + }, + }, + }, + }; + const store = createStore( + ( + currentState = state, + action: { type: string; state?: SqlLabRootState }, + ): SqlLabRootState => + action.type === 'REPLACE' && action.state ? action.state : currentState, + state, + applyMiddleware(thunk), + ); + const hook = renderHook( + ({ limit = 1000 }) => useQueryCostWarning(editor, limit), + { + initialProps: { limit: 1000 }, + wrapper: createWrapper({ useRedux: true, store }), + }, + ); + const replace = (change: (state: SqlLabRootState) => SqlLabRootState) => + act(() => { + store.dispatch({ type: 'REPLACE', state: change(store.getState()) }); + }); + const finish = async (result: object, index = 0) => + act(async () => { + resolvers[index]({ result }); + await new Promise(resolve => setTimeout(resolve, 0)); + }); + const start = () => + act(() => { + hook.result.current.start(); + }); + const executions = () => fetchMock.callHistory.calls('warning-execute'); + return { ...hook, replace, finish, start, executions, editor, store }; +} + +beforeEach(() => { + featureEnabled.mockImplementation(flag => + [ + FeatureFlag.EstimateQueryCost, + FeatureFlag.SqlLabQueryCostWarnings, + ].includes(flag), + ); +}); +afterEach(() => { + fetchMock.clearHistory().removeRoutes(); + featureEnabled.mockReset(); +}); + +test.each(['disabled', 'unsupported', 'within_threshold'])( + 'executes a captured request once for %s', + async status => { + const h = setup(); + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + unsavedQueryEditor: { + id: h.editor.id, + sql: 'SELECT 99', + selectedText: 'SELECT 2', + catalog: 'catalog', + schema: 'schema', + templateParams: '{"x":2}', + queryLimit: 25, + }, + }, + })); + h.start(); + h.start(); + expect(h.result.current.checking).toBe(true); + expect(h.result.current.hasPending()).toBe(true); + await h.finish({ ...warning, status, value: 0 }); + await waitFor(() => expect(h.executions()).toHaveLength(1)); + const checked = JSON.parse( + String(fetchMock.callHistory.calls('warning-check')[0].options.body), + ); + const executed = JSON.parse(String(h.executions()[0].options.body)); + expect(checked).toEqual({ + check_warning: true, + database_id: h.editor.dbId, + sql: 'SELECT 2', + catalog: 'catalog', + schema: 'schema', + template_params: { x: 2 }, + query_limit: 25, + }); + expect(executed).toMatchObject({ + database_id: checked.database_id, + sql: checked.sql, + catalog: checked.catalog, + schema: checked.schema, + queryLimit: checked.query_limit, + templateParams: '{"x":2}', + runAsync: false, + }); + expect(h.result.current.checking).toBe(false); + expect(h.result.current.hasPending()).toBe(false); + }, +); + +test.each(['exceeded', 'unavailable'])( + 'requires confirmation for %s and continues once', + async status => { + const h = setup(); + h.start(); + await h.finish({ + ...warning, + status, + value: status === 'unavailable' ? null : 12.5, + }); + expect(h.executions()).toHaveLength(0); + expect(h.result.current.confirmation?.status).toBe(status); + h.start(); + expect(fetchMock.callHistory.calls('warning-check')).toHaveLength(1); + act(() => { + h.result.current.continueRun(); + h.result.current.continueRun(); + }); + await waitFor(() => expect(h.executions()).toHaveLength(1)); + expect(h.result.current.confirmation).toBeUndefined(); + }, +); + +test.each(['checking', 'dialog'])( + 'cancels during %s and ignores late results', + async stage => { + const h = setup(); + h.start(); + if (stage === 'dialog') await h.finish(warning); + act(() => h.result.current.cancel()); + if (stage === 'checking') await h.finish(warning); + act(() => h.result.current.continueRun()); + expect(h.executions()).toHaveLength(0); + expect(h.result.current.checking).toBe(false); + expect(h.result.current.confirmation).toBeUndefined(); + }, +); + +test('a cancelled response cannot replace a newer attempt', async () => { + const h = setup(); + h.start(); + act(() => h.result.current.cancel()); + h.start(); + await h.finish({ ...warning, status: 'within_threshold' }); + expect(h.result.current.checking).toBe(true); + expect(h.result.current.confirmation).toBeUndefined(); + expect(h.executions()).toHaveLength(0); + await h.finish(warning, 1); + expect(h.result.current.confirmation?.status).toBe('exceeded'); + act(() => h.result.current.continueRun()); + await waitFor(() => expect(h.executions()).toHaveLength(1)); +}); + +test.each([ + ['sql', 'SELECT 3'], + ['selectedText', 'SELECT 4'], + ['dbId', 99], + ['catalog', 'other'], + ['schema', 'other'], + ['templateParams', '{"x":1}'], + ['queryLimit', 50], +])( + 'invalidates changed %s even if restored before response', + async (field, value) => { + const h = setup(); + h.start(); + const original = h.store.getState(); + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + unsavedQueryEditor: { id: h.editor.id, [field]: value }, + }, + })); + h.replace(() => original); + await h.finish(warning); + act(() => h.result.current.continueRun()); + expect(h.executions()).toHaveLength(0); + expect(h.result.current.confirmation).toBeUndefined(); + }, +); + +test.each(['async', 'rule', 'estimation', 'metric', 'tab'])( + 'invalidates %s changes', + async field => { + const h = setup(); + h.start(); + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + ...(field === 'tab' + ? { tabHistory: ['other'] } + : { + databases: { + ...state.sqlLab.databases, + [h.editor.dbId!]: { + ...state.sqlLab.databases[h.editor.dbId!], + ...(field === 'async' && { allow_run_async: true }), + ...(field === 'rule' && { + query_cost_warning: { enabled: false }, + }), + ...(field === 'estimation' && { + allows_cost_estimate: false, + }), + ...(field === 'metric' && { + engine_information: { query_cost_warning_metric: null }, + }), + }, + }, + }), + }, + })); + await h.finish({ ...warning, status: 'within_threshold' }); + expect(h.executions()).toHaveLength(0); + }, +); + +test('invalidates a default effective limit change', async () => { + const h = setup(); + h.start(); + h.rerender({ limit: 2000 }); + await h.finish(warning); + expect(h.result.current.confirmation).toBeUndefined(); +}); + +test('tab changes close an existing confirmation and invalidate its action', async () => { + const h = setup(); + h.start(); + await h.finish(warning); + const { continueRun } = h.result.current; + h.replace(state => ({ + ...state, + sqlLab: { ...state.sqlLab, tabHistory: ['other'] }, + })); + act(() => continueRun()); + expect(h.result.current.confirmation).toBeUndefined(); + expect(h.executions()).toHaveLength(0); +}); + +test('unmounted editors ignore a late response', async () => { + const h = setup(); + h.start(); + h.unmount(); + await h.finish(warning); + expect(h.executions()).toHaveLength(0); +}); + +test.each([ + {}, + { status: 'bogus' }, + { ...warning, value: '12' }, + { ...warning, value: NaN }, + { ...warning, threshold: 0 }, +])('malformed response %j never executes', async result => { + const h = setup(); + h.start(); + await h.finish(result); + expect(h.executions()).toHaveLength(0); + expect(h.result.current.checking).toBe(false); + expect(h.result.current.confirmation).toBeUndefined(); +}); + +test.each([403, 400, 500])( + 'HTTP %s remains an error and never executes', + async status => { + const h = setup(); + fetchMock.modifyRoute('warning-check', { + response: { status, body: { message: 'Check failed' } }, + }); + h.start(); + await waitFor(() => expect(h.result.current.checking).toBe(false)); + expect(h.executions()).toHaveLength(0); + expect(h.result.current.confirmation).toBeUndefined(); + }, +); + +test.each(['flag', 'estimation', 'unsupported', 'ctas', 'automatic'])( + 'known inactive %s keeps the existing execution path', + async reason => { + const h = setup(); + if (reason === 'flag') featureEnabled.mockReturnValue(false); + if (['estimation', 'rule', 'unsupported'].includes(reason)) + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + databases: { + [h.editor.dbId!]: { + ...state.sqlLab.databases[h.editor.dbId!], + ...(reason === 'estimation' && { allows_cost_estimate: false }), + ...(reason === 'rule' && { + query_cost_warning: { enabled: false }, + }), + ...(reason === 'unsupported' && { + engine_information: { query_cost_warning_metric: null }, + }), + }, + }, + }, + })); + act(() => { + h.result.current.start( + '', + reason === 'ctas', + 'TABLE', + reason === 'automatic', + ); + }); + await waitFor(() => expect(h.executions()).toHaveLength(1)); + expect(fetchMock.callHistory.calls('warning-check')).toHaveLength(0); + }, +); + +test('byte results preserve values beyond JavaScript safe integers', () => { + const value = BigInt('9007199254740993123'); + expect( + parseWarningResult({ ...warning, metric: 'bytes_processed', value }), + ).toMatchObject({ value }); + expect(() => + parseWarningResult({ + ...warning, + metric: 'bytes_processed', + value: 9007199254740992, + }), + ).toThrow(); +}); + +test('a disabled display fallback still lets the server validate the saved rule', async () => { + const h = setup(); + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + databases: { + [h.editor.dbId!]: { + ...state.sqlLab.databases[h.editor.dbId!], + query_cost_warning: { enabled: false }, + }, + }, + }, + })); + h.start(); + await h.finish({ status: 'disabled' }); + await waitFor(() => expect(h.executions()).toHaveLength(1)); + expect(fetchMock.callHistory.calls('warning-check')).toHaveLength(1); +}); + +test('the existing bigint response parser preserves large estimated bytes over HTTP', async () => { + const h = setup(); + fetchMock.modifyRoute('warning-check', { + response: { + body: '{"result":{"status":"exceeded","metric":"bytes_processed","label":"Estimated bytes processed","unit":"bytes","value":9007199254740993123,"threshold":100}}', + headers: { 'Content-Type': 'application/json' }, + }, + }); + h.start(); + await waitFor(() => + expect(h.result.current.confirmation?.value).toBe( + BigInt('9007199254740993123'), + ), + ); + expect(h.executions()).toHaveLength(0); +}); + +test('invalid template JSON is an error before an estimate or execution request', async () => { + const h = setup(); + h.replace(state => ({ + ...state, + sqlLab: { + ...state.sqlLab, + unsavedQueryEditor: { id: h.editor.id, templateParams: '{invalid' }, + }, + })); + h.start(); + await waitFor(() => expect(h.result.current.checking).toBe(false)); + expect(fetchMock.callHistory.calls('warning-check')).toHaveLength(0); + expect(h.executions()).toHaveLength(0); +}); diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.ts b/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.ts new file mode 100644 index 000000000000..3053ee1635ed --- /dev/null +++ b/superset-frontend/src/SqlLab/components/SqlEditor/useQueryCostWarning.ts @@ -0,0 +1,276 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import { useCallback, useEffect, useRef, useState } from 'react'; +import { useStore } from 'react-redux'; +import { t } from '@apache-superset/core/translation'; +import { + FeatureFlag, + getClientErrorObject, + isFeatureEnabled, + SupersetClient, +} from '@superset-ui/core'; +import type { DatabaseObject } from 'src/features/databases/types'; +import { addDangerToast } from 'src/components/MessageToasts/actions'; +import { + CtasEnum, + getSqlEditorQuery, + getUpToDateQuery, + Query, + runQuery, + setActiveSouthPaneTab, +} from 'src/SqlLab/actions/sqlLab'; +import { useAppDispatch } from 'src/SqlLab/hooks/useAppDispatch'; +import type { QueryEditor, SqlLabRootState } from 'src/SqlLab/types'; + +export type WarningEstimate = { + status: 'exceeded' | 'unavailable' | 'within_threshold'; + metric: 'bytes_processed' | 'total_cost'; + label: string; + unit: string; + value: number | bigint | null; + threshold: number; + message?: string; + accuracy?: string; +}; +type CheckResult = WarningEstimate | { status: 'disabled' | 'unsupported' }; +type Attempt = { query: Query; key: string; result?: WarningEstimate }; + +/** Reject incomplete responses rather than treating them as permission to execute. */ +export function parseWarningResult(result: unknown): CheckResult { + const invalid = () => + new Error(t('The query cost warning response was invalid.')); + if (!result || typeof result !== 'object') throw invalid(); + const value = result as Record; + if (value.status === 'disabled' || value.status === 'unsupported') { + return { status: value.status }; + } + if ( + !['exceeded', 'within_threshold', 'unavailable'].includes( + String(value.status), + ) || + !['bytes_processed', 'total_cost'].includes(String(value.metric)) || + typeof value.label !== 'string' || + !value.label.trim() || + typeof value.unit !== 'string' || + !value.unit.trim() || + typeof value.threshold !== 'number' || + !Number.isFinite(value.threshold) || + value.threshold <= 0 || + (value.message !== undefined && typeof value.message !== 'string') || + (value.accuracy !== undefined && typeof value.accuracy !== 'string') + ) + throw invalid(); + const usable = + value.metric === 'bytes_processed' + ? (typeof value.value === 'bigint' && value.value >= BigInt(0)) || + (typeof value.value === 'number' && + Number.isSafeInteger(value.value) && + value.value >= 0) + : typeof value.value === 'number' && + Number.isFinite(value.value) && + value.value >= 0; + if ( + (!usable && !(value.status === 'unavailable' && value.value === null)) || + (value.metric === 'bytes_processed' && + !Number.isSafeInteger(value.threshold)) || + (value.status === 'unavailable' && + (typeof value.message !== 'string' || !value.message.trim())) + ) + throw invalid(); + return value as WarningEstimate; +} + +/** Own one captured Run attempt; cancellation invalidates callbacks, not database jobs. */ +export default function useQueryCostWarning( + queryEditor: QueryEditor, + defaultQueryLimit: number, +) { + const store = useStore(); + const dispatch = useAppDispatch(); + const current = useRef(); + const mounted = useRef(false); + const inputs = useRef({ queryEditor, defaultQueryLimit }); + inputs.current = { queryEditor, defaultQueryLimit }; + const [pending, setPending] = useState(); + + const readInputs = useCallback(() => { + const state = store.getState(); + const props = inputs.current; + const editor = getUpToDateQuery(state, props.queryEditor); + const database: DatabaseObject | undefined = + state.sqlLab.databases[editor.dbId || '']; + const flags = [ + isFeatureEnabled(FeatureFlag.SqlLabQueryCostWarnings), + isFeatureEnabled(FeatureFlag.EstimateQueryCost), + ]; + const key = JSON.stringify({ + id: editor.id, + active: state.sqlLab.tabHistory.slice(-1)[0], + sql: editor.sql, + selectedText: editor.selectedText, + dbId: editor.dbId, + catalog: editor.catalog, + schema: editor.schema, + templateParams: editor.templateParams, + limit: editor.queryLimit || props.defaultQueryLimit, + async: database?.allow_run_async, + tab: editor.name, + tabViewId: editor.tabViewId, + immutableId: editor.immutableId, + backend: database?.backend, + estimation: database?.allows_cost_estimate, + rule: database?.query_cost_warning, + metric: database?.engine_information?.query_cost_warning_metric, + flags, + }); + return { + editor, + database, + flags, + key, + active: state.sqlLab.tabHistory.slice(-1)[0] === editor.id, + }; + }, [store]); + + const cancel = useCallback(() => { + current.current = undefined; + setPending(undefined); + }, []); + const hasPending = useCallback(() => !!current.current, []); + const isCurrent = useCallback( + (attempt: Attempt) => + mounted.current && + current.current === attempt && + readInputs().key === attempt.key, + [readInputs], + ); + useEffect(() => { + mounted.current = true; + const unsubscribe = store.subscribe(() => { + if (current.current && !isCurrent(current.current)) cancel(); + }); + return () => { + mounted.current = false; + current.current = undefined; + unsubscribe(); + }; + }, [cancel, isCurrent, store]); + useEffect(() => { + if (current.current && !isCurrent(current.current)) cancel(); + }); + + const execute = useCallback( + (query: Query) => { + dispatch(runQuery(query)); + dispatch(setActiveSouthPaneTab('Results')); + }, + [dispatch], + ); + const continueRun = useCallback(() => { + const attempt = current.current; + if (!attempt?.result || !isCurrent(attempt)) { + cancel(); + return; + } + cancel(); + execute(attempt.query); + }, [cancel, execute, isCurrent]); + + const start = useCallback( + async ( + tempTable = '', + ctas = false, + ctasMethod: string = CtasEnum.Table, + automatic = false, + ) => { + if (current.current) return; + const { editor, database, flags, key, active } = readInputs(); + if (!database || (!active && !automatic)) return; + const query = getSqlEditorQuery( + { id: database.id, allow_run_async: Boolean(database.allow_run_async) }, + editor, + inputs.current.defaultQueryLimit, + tempTable, + ctas, + ctasMethod, + ); + if ( + ctas || + automatic || + flags.some(flag => !flag) || + // The display rule masks invalid settings; only the server can disable it. + database.allows_cost_estimate === false || + database.engine_information?.query_cost_warning_metric === null + ) { + execute(query); + return; + } + const attempt: Attempt = { query, key }; + current.current = attempt; + setPending(attempt); + try { + const { json } = await SupersetClient.post({ + endpoint: '/api/v1/sqllab/estimate/', + jsonPayload: { + check_warning: true, + database_id: query.dbId, + sql: query.sql, + catalog: query.catalog, + schema: query.schema, + template_params: JSON.parse(query.templateParams || '{}'), + query_limit: query.queryLimit, + }, + parseMethod: 'json-bigint', + }); + if (!isCurrent(attempt)) return; + const result = parseWarningResult(json.result); + if (result.status === 'exceeded' || result.status === 'unavailable') { + const confirmation = { ...attempt, result }; + current.current = confirmation; + setPending(confirmation); + } else { + cancel(); + execute(query); + } + } catch (error) { + if (!isCurrent(attempt)) return; + const clientError = await getClientErrorObject(error); + if (!isCurrent(attempt)) return; + cancel(); + dispatch( + addDangerToast( + clientError.error || + clientError.message || + clientError.statusText || + t('Query cost check failed.'), + ), + ); + } + }, + [cancel, dispatch, execute, isCurrent, readInputs], + ); + return { + start, + hasPending, + cancel, + continueRun, + checking: !!pending && !pending.result, + confirmation: pending?.result, + }; +} From 56aa31285e36e9db9594b2912228f1d4b417242d Mon Sep 17 00:00:00 2001 From: Alankrit Verma Date: Sun, 18 Oct 2026 18:11:23 -0400 Subject: [PATCH 11/58] fix(sqllab): capture current statement without changing selection --- docs/docs/using-superset/exploring-data.mdx | 3 +- .../SqlEditor/SqlEditor.entryPoints.test.tsx | 290 ++++++++++++++++++ .../src/SqlLab/components/SqlEditor/index.tsx | 37 +-- .../SqlEditor/useQueryCostWarning.test.tsx | 2 - .../SqlEditor/useQueryCostWarning.ts | 5 +- 5 files changed, 313 insertions(+), 24 deletions(-) create mode 100644 superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.entryPoints.test.tsx diff --git a/docs/docs/using-superset/exploring-data.mdx b/docs/docs/using-superset/exploring-data.mdx index 0a99ff07f06f..0816cb3a671b 100644 --- a/docs/docs/using-superset/exploring-data.mdx +++ b/docs/docs/using-superset/exploring-data.mdx @@ -468,7 +468,8 @@ database or network errors remain errors. Omitting `check_warning` preserves the manual estimator response. When you deliberately select **Run** (including Run selection or its keyboard -shortcut), SQL Lab checks the captured request before execution. While checking, +shortcut), SQL Lab checks the captured request before execution. The current-statement +shortcut uses the existing statement range without changing your selection. While checking, **Cancel** in the Run area prevents that attempt from executing; it does not forcibly cancel a database dry run. An estimate above the threshold offers **Cancel** and **Run anyway**. An unavailable estimate explains why and offers diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.entryPoints.test.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.entryPoints.test.tsx new file mode 100644 index 000000000000..4c8948f8f774 --- /dev/null +++ b/superset-frontend/src/SqlLab/components/SqlEditor/SqlEditor.entryPoints.test.tsx @@ -0,0 +1,290 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import type { editors } from '@apache-superset/core'; +import { isFeatureEnabled, FeatureFlag } from '@superset-ui/core'; +import { useSelector } from 'react-redux'; +import fetchMock from 'fetch-mock'; +import { + act, + fireEvent, + render, + screen, + waitFor, +} from 'spec/helpers/testing-library'; +import reducers from 'spec/helpers/reducerIndex'; +import { setupStore } from 'src/views/store'; +import { initialState, defaultQueryEditor } from 'src/SqlLab/fixtures'; +import { queryEditorSetSelectedText } from 'src/SqlLab/actions/sqlLab'; +import type { SqlLabRootState } from 'src/SqlLab/types'; +import SqlEditor from '.'; + +jest.mock('@superset-ui/core', () => ({ + ...jest.requireActual('@superset-ui/core'), + isFeatureEnabled: jest.fn(), +})); +jest.mock( + 'react-virtualized-auto-sizer', + () => + ({ + children, + }: { + children: (size: { height: number }) => React.ReactNode; + }) => + children({ height: 500 }), +); +jest.mock('src/SqlLab/components/ResultSet', () => () =>
); +jest.mock('src/components/DatabaseSelector', () => ({ + DatabaseSelector: () =>
, +})); + +let mockCursor = { line: 0, column: 0 }; +const mockSetSelection = jest.fn(); +const mockMoveCursor = jest.fn(); +const mockScrollToLine = jest.fn(); +jest.mock('src/core/editors', () => ({ + ...jest.requireActual('src/core/editors'), + EditorHost: ({ value, hotkeys, onChange }: editors.EditorProps) => { + const handle: editors.EditorHandle = { + getValue: () => value, + getCursorPosition: () => mockCursor, + setSelection: mockSetSelection, + moveCursorToPosition: mockMoveCursor, + scrollToLine: mockScrollToLine, + focus: jest.fn(), + setValue: jest.fn(), + getSelections: () => [], + getSelectedText: () => '', + insertText: jest.fn(), + executeCommand: jest.fn(), + setAnnotations: jest.fn(), + clearAnnotations: jest.fn(), + resize: jest.fn(), + registerCompletionProvider: () => ({ dispose: jest.fn() }), + onDidChangeContent: () => ({ dispose: jest.fn() }), + }; + return ( +