diff --git a/superset-core/src/superset_core/tasks/models.py b/superset-core/src/superset_core/tasks/models.py index 075f2e0c899..e1db14d0fb1 100644 --- a/superset-core/src/superset_core/tasks/models.py +++ b/superset-core/src/superset_core/tasks/models.py @@ -145,7 +145,9 @@ class TaskSubscriber(CoreModel): This model tracks task subscriptions for multi-user shared tasks. When a user schedules a shared task with the same parameters as an existing task, - they are subscribed to that task instead of creating a duplicate. + they are subscribed to that task instead of creating a duplicate. A subscriber + is identified by exactly one of ``user_id`` (authenticated) or ``guest_key`` + (an embedded guest, which has no ``ab_user`` row). """ __abstract__ = True @@ -153,7 +155,8 @@ class TaskSubscriber(CoreModel): # Type hints for expected attributes (no actual field definitions) id: int task_id: int - user_id: int + user_id: int | None + guest_key: str | None subscribed_at: datetime # Audit fields from AuditMixinNullable diff --git a/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.test.tsx b/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.test.tsx index f66ad3a516f..b68f08200b9 100644 --- a/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.test.tsx +++ b/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.test.tsx @@ -714,10 +714,7 @@ test('should refetch when mixing renderTrigger string control with non-renderTri test('resolves async (202) responses via the injected handleAsyncChartData hook', async () => { const asyncJob = { - channel_id: 'c1', - job_id: 'j1', - status: 'running', - result_url: '/api/v1/chart/data/abc', + task_ids: ['task-1', 'task-2'], }; mockChartClient.client.post.mockResolvedValue({ response: { status: 202 } as Response, @@ -738,10 +735,12 @@ test('resolves async (202) responses via the injected handleAsyncChartData hook' await waitFor(() => { expect(handleAsyncChartData).toHaveBeenCalledTimes(1); }); - // Delegates the raw response + job metadata (and abort signal) + // Delegates the raw response + async job (task_ids), a refetch thunk, and the + // abort signal. expect(handleAsyncChartData).toHaveBeenCalledWith( { status: 202 }, asyncJob, + expect.any(Function), expect.any(AbortSignal), ); // Chart renders once the async data resolves @@ -978,7 +977,7 @@ test('passes an abort signal to the async handler and aborts it on unmount', asy response: { status: 202 } as Response, json: { job_id: 'j', channel_id: 'c' }, }); - // Typed with a rest param so mock.calls is indexable (the 3rd arg is the signal) + // Typed with a rest param so mock.calls is indexable (the 4th arg is the signal) const handleAsyncChartData = jest.fn( (..._args: unknown[]) => new Promise(() => {}), // never resolves ); @@ -994,7 +993,7 @@ test('passes an abort signal to the async handler and aborts it on unmount', asy await waitFor(() => { expect(handleAsyncChartData).toHaveBeenCalledTimes(1); }); - const signal = handleAsyncChartData.mock.calls[0][2] as AbortSignal; + const signal = handleAsyncChartData.mock.calls[0][3] as AbortSignal; expect(signal).toBeInstanceOf(AbortSignal); expect(signal.aborted).toBe(false); diff --git a/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.tsx b/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.tsx index e56b7820911..f7faeaef640 100644 --- a/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.tsx +++ b/superset-frontend/packages/superset-ui-core/src/chart/components/StatefulChart.tsx @@ -319,12 +319,12 @@ export default function StatefulChart(props: StatefulChartProps) { let responseData: QueryData[]; if (rawResponse?.status === 202) { - // With GLOBAL_ASYNC_QUERIES the query is dispatched to a Celery worker - // and the 202 body is job metadata (channel_id, job_id, result_url), - // not chart data. Delegate to the injected handler, which polls the - // async event channel and resolves the cached results. Without a - // handler we fail loudly rather than rendering the job metadata as if - // it were an (empty) result set. + // With GLOBAL_ASYNC_QUERIES the query runs as one GTF task per + // QueryObject and the 202 body is the async job ({task_ids}), not chart + // data. Delegate to the injected handler, which polls task statuses and, + // once they succeed, calls `refetch` to re-issue this request and read + // the now-cached results. Without a handler we fail loudly rather than + // rendering the job metadata as if it were an (empty) result set. if (!hooks?.handleAsyncChartData) { throw new Error( 'Received an async chart data response (HTTP 202) but no async ' + @@ -332,10 +332,24 @@ export default function StatefulChart(props: StatefulChartProps) { 'the async handler or disable GLOBAL_ASYNC_QUERIES for this chart.', ); } + // Re-issue from the warm per-query cache (force off) and extract rows. + const refetch = async (): Promise => { + const cached = await chartClientRef.current!.client.post({ + ...requestConfig, + jsonPayload: queryContext, + }); + const cachedRows = ( + Array.isArray(cached.json) ? cached.json : [cached.json] + ) as JsonObject[]; + return ( + cachedRows[0]?.result ? cachedRows[0].result : cachedRows + ) as QueryData[]; + }; responseData = ensureIsArray( await hooks.handleAsyncChartData( rawResponse, clientResponse.json as JsonObject, + refetch, controller.signal, ), ); diff --git a/superset-frontend/packages/superset-ui-core/src/chart/models/ChartProps.ts b/superset-frontend/packages/superset-ui-core/src/chart/models/ChartProps.ts index 638b5888a90..bfd48d6b731 100644 --- a/superset-frontend/packages/superset-ui-core/src/chart/models/ChartProps.ts +++ b/superset-frontend/packages/superset-ui-core/src/chart/models/ChartProps.ts @@ -75,6 +75,7 @@ type Hooks = { handleAsyncChartData?: ( response: Response, json: JsonObject, + refetch?: () => Promise, signal?: AbortSignal, ) => Promise | QueryData[]; } & PlainObject; diff --git a/superset-frontend/src/components/Chart/ChartRenderer.tsx b/superset-frontend/src/components/Chart/ChartRenderer.tsx index 59cbb1779c3..3672e0662f3 100644 --- a/superset-frontend/src/components/Chart/ChartRenderer.tsx +++ b/superset-frontend/src/components/Chart/ChartRenderer.tsx @@ -169,6 +169,7 @@ interface ChartHooks { handleAsyncChartData?: ( response: Response, json: JsonObject, + refetch?: () => Promise, signal?: AbortSignal, ) => Promise | QueryData[]; } diff --git a/superset-frontend/src/components/Chart/DrillBy/DrillByModal.tsx b/superset-frontend/src/components/Chart/DrillBy/DrillByModal.tsx index fc490a2d257..485c9a7cff5 100644 --- a/superset-frontend/src/components/Chart/DrillBy/DrillByModal.tsx +++ b/superset-frontend/src/components/Chart/DrillBy/DrillByModal.tsx @@ -402,10 +402,18 @@ export default function DrillByModal({ setChartDataResult(undefined); setIsChartDataLoading(true); - getChartDataRequest({ - formData: drilledFormData, - }) - .then(({ response, json }) => handleChartDataResponse(response, json)) + const requestDrillData = () => + getChartDataRequest({ + formData: drilledFormData, + }); + requestDrillData() + .then(({ response, json }) => + handleChartDataResponse(response, json, () => + requestDrillData().then(({ response: r, json: j }) => + handleChartDataResponse(r, j), + ), + ), + ) .then(queriesResponse => { setChartDataResult(queriesResponse); }) @@ -491,10 +499,18 @@ export default function DrillByModal({ if (drilledFormData) { setIsChartDataLoading(true); setChartDataResult(undefined); - getChartDataRequest({ - formData: drilledFormData, - }) - .then(({ response, json }) => handleChartDataResponse(response, json)) + const requestDrillData = () => + getChartDataRequest({ + formData: drilledFormData, + }); + requestDrillData() + .then(({ response, json }) => + handleChartDataResponse(response, json, () => + requestDrillData().then(({ response: r, json: j }) => + handleChartDataResponse(r, j), + ), + ), + ) .then(queriesResponse => { setChartDataResult(queriesResponse); }) diff --git a/superset-frontend/src/components/Chart/chartAction.ts b/superset-frontend/src/components/Chart/chartAction.ts index 7ead1472e96..f60c13b41ad 100644 --- a/superset-frontend/src/components/Chart/chartAction.ts +++ b/superset-frontend/src/components/Chart/chartAction.ts @@ -48,7 +48,7 @@ import { logEvent } from 'src/logger/actions'; import { Logger, LOG_ACTIONS_LOAD_CHART } from 'src/logger/LogUtils'; import { allowCrossDomain as domainShardingEnabled } from 'src/utils/hostNamesConfig'; import { updateDataMask } from 'src/dataMask/actions'; -import { waitForAsyncData } from 'src/middleware/asyncEvent'; +import { AsyncJob, waitForAsyncData } from 'src/middleware/asyncEvent'; import { ensureAppRoot } from 'src/utils/navigationUtils'; import { safeStringify } from 'src/utils/safeStringify'; import { extendedDayjs } from '@superset-ui/core/utils/dates'; @@ -639,6 +639,7 @@ export function addChart( export function handleChartDataResponse( response: Response, json: { result: QueryData[] }, + refetch?: () => Promise, signal?: AbortSignal, ): Promise | QueryData[] { if (isFeatureEnabled(FeatureFlag.GlobalAsyncQueries)) { @@ -648,16 +649,19 @@ export function handleChartDataResponse( case 200: // Query results returned synchronously, meaning query was already cached. return Promise.resolve(result); - case 202: - // Query is running asynchronously and we must await the results. - // When status is 202, result contains async event data (job_id, channel_id, etc.) - // which differs from QueryData. We cast through unknown to handle this safely. - // The optional signal lets a caller abort the wait (Stop pressed, chart - // superseded or unmounted), cancelling the job and avoiding leaked listeners. - return waitForAsyncData( - result as unknown as Parameters[0], - signal, - ) as Promise; + case 202: { + // Query is running asynchronously as one GTF task per QueryObject. The + // 202 body is the async job ({task_ids}); await every task, then re-issue + // this request via `refetch` to read the now-cached results. The optional + // signal lets a caller abort the wait (Stop pressed, chart superseded or + // unmounted), cancelling the outstanding tasks. + if (!refetch) { + throw new Error( + 'Async chart-data response (202) received without a refetch handler', + ); + } + return waitForAsyncData(json as unknown as AsyncJob, refetch, signal); + } default: throw new Error( `Received unexpected response status (${response.status}) while fetching chart data`, @@ -705,19 +709,33 @@ export function exploreJSON( setTimeout(() => prevController.abort(), 0); } - const chartDataRequest = getChartDataRequest({ - setDataMask, - formData, - resultFormat: 'json', - resultType: 'full', - force, - requestParams, - ownState, - }); + // Re-issue the chart-data request. On the async path this runs after every + // query task has succeeded, so `force` is dropped — the per-query DATA cache + // is warm and this returns synchronously (200) from cache. + const requestChartData = (fromCache = false) => + getChartDataRequest({ + setDataMask, + formData, + resultFormat: 'json', + resultType: 'full', + force: fromCache ? false : force, + requestParams, + ownState, + }); + + const chartDataRequest = requestChartData(); const chartDataRequestCaught = chartDataRequest .then(({ response, json }) => - handleChartDataResponse(response, json, controller.signal), + handleChartDataResponse( + response, + json, + () => + requestChartData(true).then(({ response: r, json: j }) => + handleChartDataResponse(r, j), + ) as Promise, + controller.signal, + ), ) .then(queriesResponse => { // Drop stale responses: if this request was aborted (Stop, or a newer diff --git a/superset-frontend/src/components/Chart/chartActions.test.ts b/superset-frontend/src/components/Chart/chartActions.test.ts index c1d0781d056..42197977671 100644 --- a/superset-frontend/src/components/Chart/chartActions.test.ts +++ b/superset-frontend/src/components/Chart/chartActions.test.ts @@ -26,6 +26,7 @@ import { getChartBuildQueryRegistry, QueryFormData, JsonObject, + QueryData, AnnotationLayer, AnnotationType, AnnotationSourceType, @@ -153,7 +154,10 @@ describe('chart actions', () => { ); waitForAsyncDataStub = jest .spyOn(asyncEvent, 'waitForAsyncData') - .mockImplementation((data: unknown) => Promise.resolve(data)); + // New contract: resolve by invoking the caller-provided refetch thunk. + .mockImplementation((_job: unknown, refetch: () => Promise) => + refetch(), + ); }); test('should drop stale success dispatches when a newer controller has replaced ours in state', async () => { @@ -402,14 +406,20 @@ describe('chart actions', () => { ).featureFlags = { [FeatureFlag.GlobalAsyncQueries]: true, }; + // On 202 the body is the async job ({task_ids}); once the tasks resolve + // (stubbed waitForAsyncData invokes the refetch), the re-request returns + // the cached data. + const refetch = jest + .fn() + .mockResolvedValue([1, 2, 3] as unknown as QueryData[]); const result = await handleChartDataResponse( { status: 202 } as Response, { - result: [ - 1, 2, 3, - ] as unknown as actions.ChartDataRequestResponse['json']['result'], - }, + task_ids: ['task-1'], + } as unknown as actions.ChartDataRequestResponse['json'], + refetch, ); + expect(refetch).toHaveBeenCalledTimes(1); expect(result).toEqual([1, 2, 3]); }); diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FilterBar/FilterControls/FilterValue.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FilterBar/FilterControls/FilterValue.tsx index 359fdcff50f..3a77df06c95 100644 --- a/superset-frontend/src/dashboard/components/nativeFilters/FilterBar/FilterControls/FilterValue.tsx +++ b/superset-frontend/src/dashboard/components/nativeFilters/FilterBar/FilterControls/FilterValue.tsx @@ -49,7 +49,7 @@ import { isEqual, isEqualWith } from 'lodash-es'; import { getChartDataRequest } from 'src/components/Chart/chartAction'; import { ErrorAlert, ErrorMessageWithStackTrace } from 'src/components'; import { Loading, Constants, Flex } from '@superset-ui/core/components'; -import { waitForAsyncData } from 'src/middleware/asyncEvent'; +import { waitForAsyncData, AsyncJob } from 'src/middleware/asyncEvent'; import { FilterBarOrientation, RootState } from 'src/dashboard/types'; import { onFiltersRefreshSuccess, @@ -278,11 +278,13 @@ const FilterValue: FC = ({ return; } setIsRefreshing(true); - getChartDataRequest({ - formData: newFormData, - force: shouldRefresh, - ownState: filterOwnState, - }) + const requestFilterData = (fromCache = false) => + getChartDataRequest({ + formData: newFormData, + force: fromCache ? false : shouldRefresh, + ownState: filterOwnState, + }); + requestFilterData() .then(({ response, json }) => { if (isFeatureEnabled(FeatureFlag.GlobalAsyncQueries)) { // deal with getChartDataRequest transforming the response data @@ -292,7 +294,14 @@ const FilterValue: FC = ({ setError(undefined); handleFilterLoadFinish(); } else if (response.status === 202) { - waitForAsyncData(result as Parameters[0]) + // Await the query tasks, then re-issue the request to read the + // now-cached results. + waitForAsyncData(json as unknown as AsyncJob, () => + requestFilterData(true).then( + ({ json: cachedJson }) => + cachedJson.result as ChartDataResponseResult[], + ), + ) .then((asyncResult: ChartDataResponseResult[]) => { setState(asyncResult); setError(undefined); diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx index 9c7ff4dbbcb..8f0164dedcf 100644 --- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx +++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx @@ -85,7 +85,7 @@ import { import DateFilterControl from 'src/explore/components/controls/DateFilterControl'; import AdhocFilterControl from 'src/explore/components/controls/FilterControl/AdhocFilterControl'; import type AdhocFilterClass from 'src/explore/components/controls/FilterControl/AdhocFilter'; -import { waitForAsyncData } from 'src/middleware/asyncEvent'; +import { waitForAsyncData, AsyncJob } from 'src/middleware/asyncEvent'; import { SingleValueType } from 'src/filters/components/Range/SingleValueType'; import { RangeDisplayMode } from 'src/filters/components/Range/types'; import { @@ -526,10 +526,12 @@ const FiltersConfigForm = ( defaultValueQueriesData: null, isDataDirty: false, }); - getChartDataRequest({ - formData, - force, - }) + const requestDefaultValues = (fromCache = false) => + getChartDataRequest({ + formData, + force: fromCache ? false : force, + }); + requestDefaultValues() .then(({ response, json }) => { if (isFeatureEnabled(FeatureFlag.GlobalAsyncQueries)) { // deal with getChartDataRequest transforming the response data @@ -540,7 +542,14 @@ const FiltersConfigForm = ( defaultValueQueriesData: [result as ChartDataResponseResult], }); } else if (response.status === 202) { - waitForAsyncData(result as Parameters[0]) + // Await the query tasks, then re-issue the request to read the + // now-cached results. + waitForAsyncData(json as unknown as AsyncJob, () => + requestDefaultValues(true).then( + ({ json: cachedJson }) => + cachedJson.result as ChartDataResponseResult[], + ), + ) .then((asyncResult: ChartDataResponseResult[]) => { setNativeFilterFieldValuesWrapper({ defaultValueQueriesData: asyncResult, diff --git a/superset-frontend/src/features/versionHistory/ChartVersionPreview.tsx b/superset-frontend/src/features/versionHistory/ChartVersionPreview.tsx index 8d7c24d40dd..ef09061e848 100644 --- a/superset-frontend/src/features/versionHistory/ChartVersionPreview.tsx +++ b/superset-frontend/src/features/versionHistory/ChartVersionPreview.tsx @@ -118,6 +118,9 @@ export default function ChartVersionPreview() { } fetchIdRef.current += 1; const fetchId = fetchIdRef.current; + // Abort the async chart-data wait (stop polling + cancel the GTF tasks) when + // the preview is superseded or unmounts, not just suppress the state update. + const controller = new AbortController(); setIsLoading(true); setError(null); setQueriesData(null); @@ -175,10 +178,20 @@ export default function ChartVersionPreview() { datasourceId, datasourceType, ); - const { response, json } = await getChartDataRequest({ - formData: previewFormData, - }); - const result = await handleChartDataResponse(response, json); + const requestPreviewData = () => + getChartDataRequest({ + formData: previewFormData, + }); + const { response, json } = await requestPreviewData(); + const result = await handleChartDataResponse( + response, + json, + () => + requestPreviewData().then(({ response: r, json: j }) => + handleChartDataResponse(r, j), + ), + controller.signal, + ); if (fetchId !== fetchIdRef.current) { return; } @@ -216,6 +229,7 @@ export default function ChartVersionPreview() { // afterwards. return () => { fetchIdRef.current += 1; + controller.abort(); }; }, [dispatch, entityUuid, versionUuid]); diff --git a/superset-frontend/src/middleware/asyncEvent.test.ts b/superset-frontend/src/middleware/asyncEvent.test.ts index ab7dde9b968..2de71a75a52 100644 --- a/superset-frontend/src/middleware/asyncEvent.test.ts +++ b/superset-frontend/src/middleware/asyncEvent.test.ts @@ -17,12 +17,7 @@ * under the License. */ import fetchMock from 'fetch-mock'; -import WS from 'jest-websocket-mock'; -import { - parseErrorJson, - isFeatureEnabled, - SupersetClient, -} from '@superset-ui/core'; +import { isFeatureEnabled } from '@superset-ui/core'; import * as asyncEvent from 'src/middleware/asyncEvent'; jest.mock('@superset-ui/core', () => ({ @@ -32,594 +27,133 @@ jest.mock('@superset-ui/core', () => ({ const mockedIsFeatureEnabled = isFeatureEnabled as jest.Mock; -// eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks -describe('asyncEvent middleware', () => { - const asyncPendingEvent = { - status: 'pending', - result_url: null, - job_id: 'foo123', - channel_id: '999', - errors: [], - }; - const asyncDoneEvent = { - id: '1518951480106-0', - status: 'done', - result_url: '/api/v1/chart/data/cache-key-1', - job_id: 'foo123', - channel_id: '999', - errors: [], - }; - const asyncErrorEvent = { - id: '1518951480107-0', - status: 'error', - result_url: null, - job_id: 'foo123', - channel_id: '999', - errors: [{ message: "Error: relation 'foo' does not exist" }], - }; - const chartData = { - result: [ - { - cache_key: '199f01f81f99c98693694821e4458111', - cached_dttm: null, - cache_timeout: 86400, - annotation_data: {}, - error: null, - is_cached: false, - query: - 'SELECT product_line AS product_line,\n sum(sales) AS "(Sales)"\nFROM cleaned_sales_data\nGROUP BY product_line\nLIMIT 50000', - status: 'success', - stacktrace: null, - rowcount: 7, - colnames: ['product_line', '(Sales)'], - coltypes: [1, 0], - data: [ - { - product_line: 'Classic Cars', - '(Sales)': 3919615.66, - }, - ], - applied_filters: [ - { - column: '__time_range', - }, - ], - rejected_filters: [], - }, - ], - }; +const STATUS_CHANGES_ENDPOINT = 'glob:*/api/v1/task/status_changes*'; +const CANCEL_ENDPOINT = 'glob:*/api/v1/task/*/cancel'; - const EVENTS_ENDPOINT = 'glob:*/api/v1/async_event/*'; - const CACHED_DATA_ENDPOINT = 'glob:*/api/v1/chart/data/*'; +const config = { GLOBAL_ASYNC_QUERIES_POLLING_DELAY: 20 }; - beforeEach(async () => { - mockedIsFeatureEnabled.mockImplementation( - featureFlag => featureFlag === 'GLOBAL_ASYNC_QUERIES', - ); - }); - - afterEach(() => { - fetchMock.clearHistory().removeRoutes(); - mockedIsFeatureEnabled.mockRestore(); - }); - - afterAll(() => fetchMock.clearHistory().removeRoutes()); - - // eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks - describe('polling transport', () => { - const config = { - GLOBAL_ASYNC_QUERIES_TRANSPORT: 'polling', - GLOBAL_ASYNC_QUERIES_POLLING_DELAY: 50, - GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL: '', - }; - - beforeEach(async () => { - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [asyncDoneEvent] }, - }); - fetchMock.get(CACHED_DATA_ENDPOINT, { - status: 200, - body: { result: chartData }, - }); - asyncEvent.init(config); - }); - - test('resolves with chart data on event done status', async () => { - const actualResolved = - await asyncEvent.waitForAsyncData(asyncPendingEvent); - expect(actualResolved).toEqual([chartData]); - - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1); - }); - - test('rejects with an AbortError and cancels the job when the signal aborts', async () => { - const CANCEL_ENDPOINT = 'glob:*/api/v1/async_event/*/cancel'; - fetchMock.post(CANCEL_ENDPOINT, { status: 200, body: {} }); - - const controller = new AbortController(); - const promise = asyncEvent.waitForAsyncData( - asyncPendingEvent, - controller.signal, - ); - controller.abort(); - - let error: any = null; - try { - await promise; - } catch (err) { - error = err; - } - expect(error?.name).toBe('AbortError'); - // The cancel POST is fire-and-forget; let its microtask flush. - await new Promise(resolve => setTimeout(resolve, 0)); - expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1); - }); - - test('rejects immediately when given an already-aborted signal', async () => { - const CANCEL_ENDPOINT = 'glob:*/api/v1/async_event/*/cancel'; - fetchMock.post(CANCEL_ENDPOINT, { status: 200, body: {} }); - - const controller = new AbortController(); - controller.abort(); - - await expect( - asyncEvent.waitForAsyncData(asyncPendingEvent, controller.signal), - ).rejects.toMatchObject({ name: 'AbortError' }); - // The cancel POST is fire-and-forget; let its microtask flush. - await new Promise(resolve => setTimeout(resolve, 0)); - expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1); - }); - - test('rejects on event error status', async () => { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [asyncErrorEvent] }, - }); - const errorResponse = parseErrorJson(asyncErrorEvent); - let error: any = null; - try { - await asyncEvent.waitForAsyncData(asyncPendingEvent); - } catch (err) { - error = err; - } finally { - expect(error).toEqual(errorResponse); - } - - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(0); - }); - - test('rejects on cached data fetch error', async () => { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [asyncDoneEvent] }, - }); - fetchMock.get(CACHED_DATA_ENDPOINT, { - status: 400, - }); - - let error = ''; - try { - await asyncEvent.waitForAsyncData(asyncPendingEvent); - } catch (err) { - [{ error }] = err; - } finally { - expect(error).toEqual('Bad request'); - } - - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1); - }); - - test('backs off exponentially when polling requests keep failing', async () => { - // stop the real-timer polling loop started by beforeEach before - // switching to fake timers, so all polls run on the fake clock - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - jest.useFakeTimers(); - try { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { status: 403 }); - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - - // first poll fires after the configured delay and fails - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - // next poll is delayed by 2x the configured delay, so nothing yet - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(2); - - // after the second failure the delay grows to 4x - await jest.advanceTimersByTimeAsync( - 3 * config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(2); - - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(3); - } finally { - jest.useRealTimers(); - } - }); - - test('resumes the configured polling delay after a successful poll', async () => { - // stop the real-timer polling loop started by beforeEach before - // switching to fake timers, so all polls run on the fake clock - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - jest.useFakeTimers(); - try { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { status: 403 }); - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - - // two failed polls: 1x delay, then 2x delay - await jest.advanceTimersByTimeAsync( - 3 * config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(2); - - // subsequent polls succeed, resetting the backoff - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [] }, - }); - - // third poll fires 4x delay after the second failure and succeeds - await jest.advanceTimersByTimeAsync( - 4 * config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - // polling is back to the configured delay - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(2); - } finally { - jest.useRealTimers(); - } - }); - - test('caps the polling backoff delay at 60 seconds', async () => { - const MAX_ERROR_POLLING_DELAY_MS = 60000; - // stop the real-timer polling loop started by beforeEach before - // switching to fake timers, so all polls run on the fake clock - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - jest.useFakeTimers(); - try { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { status: 403 }); - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - - // first poll fires after the configured delay and fails - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - // walk the uncapped backoff: after failure N the next delay is - // 2^N times the configured delay, which stays below the cap through - // failure 10 (50ms * 2^10 = 51.2s) - for (let failures = 1; failures <= 10; failures += 1) { - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY * 2 ** failures, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength( - failures + 1, - ); - } - - // after failure 11 the uncapped delay would be 102.4s, so the cap - // takes over: no poll just before the 60s mark... - await jest.advanceTimersByTimeAsync(MAX_ERROR_POLLING_DELAY_MS - 1); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(11); - - // ...and the next poll fires exactly at 60s - await jest.advanceTimersByTimeAsync(1); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(12); - - // additional failures remain capped at 60s - await jest.advanceTimersByTimeAsync(MAX_ERROR_POLLING_DELAY_MS); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(13); - } finally { - jest.useRealTimers(); - } - }); - - test('does not start a second loop when re-initialized during an in-flight poll', async () => { - // stop the real-timer polling loop started by beforeEach before - // switching to fake timers, so all polls run on the fake clock - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - jest.useFakeTimers(); - try { - fetchMock.clearHistory().removeRoutes(); - let resolveFetch: (response: any) => void = () => {}; - fetchMock.get( - EVENTS_ENDPOINT, - new Promise(resolve => { - resolveFetch = resolve; - }), - ); - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - - // first poll fires and stays in-flight - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - // re-init while that fetch is pending, then let it resolve; the - // stale invocation must not schedule a second loop - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - resolveFetch({ status: 200, body: { result: [] } }); - await jest.advanceTimersByTimeAsync(0); - - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [] }, - }); - - // exactly one poll per delay from here on - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(2); - } finally { - jest.useRealTimers(); - } - }); - - test('does not resume polling when re-initialized with the feature disabled during an in-flight poll', async () => { - // stop the real-timer polling loop started by beforeEach before - // switching to fake timers, so all polls run on the fake clock - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - jest.useFakeTimers(); - try { - fetchMock.clearHistory().removeRoutes(); - let resolveFetch: (response: any) => void = () => {}; - fetchMock.get( - EVENTS_ENDPOINT, - new Promise(resolve => { - resolveFetch = resolve; - }), - ); - asyncEvent.init(config); - asyncEvent.waitForAsyncData(asyncPendingEvent).catch(() => {}); - - // first poll fires and stays in-flight - await jest.advanceTimersByTimeAsync( - config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(1); - - // disable the feature and re-init while the fetch is pending; the - // stale invocation must not restart the stopped loop when it resumes - mockedIsFeatureEnabled.mockReturnValueOnce(false); - asyncEvent.init(config); - resolveFetch({ status: 200, body: { result: [] } }); - await jest.advanceTimersByTimeAsync(0); - - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [] }, - }); - - await jest.advanceTimersByTimeAsync( - 10 * config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY, - ); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(0); - } finally { - jest.useRealTimers(); - } - }); - - // Regression guard for the motivating CodeQL case: a job_id that collides - // with a built-in Object property (e.g. "__proto__"/"constructor") must be - // routed through the Map-based registries without triggering prototype - // pollution or losing the listener to a prototype-bearing lookup. - test.each(['__proto__', 'constructor', 'prototype', 'hasOwnProperty'])( - 'resolves listeners keyed by reserved job_id "%s"', - async jobId => { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [{ ...asyncDoneEvent, job_id: jobId }] }, - }); - fetchMock.get(CACHED_DATA_ENDPOINT, { - status: 200, - body: { result: chartData }, - }); - - const actualResolved = await asyncEvent.waitForAsyncData({ - ...asyncPendingEvent, - job_id: jobId, - }); - expect(actualResolved).toEqual([chartData]); - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength( - 1, - ); - }, - ); - }); - - // eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks - describe('ws transport', () => { - let wsServer: WS; - const config = { - GLOBAL_ASYNC_QUERIES_TRANSPORT: 'ws', - GLOBAL_ASYNC_QUERIES_POLLING_DELAY: 50, - GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL: 'ws://127.0.0.1:8080/', - }; - - beforeEach(async () => { - fetchMock.get(EVENTS_ENDPOINT, { - status: 200, - body: { result: [asyncDoneEvent] }, - }); - fetchMock.get(CACHED_DATA_ENDPOINT, { - status: 200, - body: { result: chartData }, - }); - - wsServer = new WS(config.GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL); - asyncEvent.init(config); - }); - - afterEach(() => { - WS.clean(); - }); - - test('resolves with chart data on event done status', async () => { - await wsServer.connected; - - const promise = asyncEvent.waitForAsyncData(asyncPendingEvent); - - wsServer.send(JSON.stringify(asyncDoneEvent)); - - await expect(promise).resolves.toEqual([chartData]); - - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(0); - }); - - test('rejects on event error status', async () => { - await wsServer.connected; - - const promise = asyncEvent.waitForAsyncData(asyncPendingEvent); - - wsServer.send(JSON.stringify(asyncErrorEvent)); - - const errorResponse = parseErrorJson(asyncErrorEvent); - - await expect(promise).rejects.toEqual(errorResponse); - - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(0); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(0); - }); - - test('rejects on cached data fetch error', async () => { - fetchMock.clearHistory().removeRoutes(); - fetchMock.get(CACHED_DATA_ENDPOINT, { - status: 400, - }); - - await wsServer.connected; - - const promise = asyncEvent.waitForAsyncData(asyncPendingEvent); - - wsServer.send(JSON.stringify(asyncDoneEvent)); - - let error = ''; - try { - await promise; - } catch (err) { - [{ error }] = err; - } finally { - expect(error).toEqual('Bad request'); - } - - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(0); - }); - - test('resolves when events are received before listener', async () => { - await wsServer.connected; - - wsServer.send(JSON.stringify(asyncDoneEvent)); - - const promise = asyncEvent.waitForAsyncData(asyncPendingEvent); - await expect(promise).resolves.toEqual([chartData]); - - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(1); - expect(fetchMock.callHistory.calls(EVENTS_ENDPOINT)).toHaveLength(0); - }); - - test('rejects with AbortError and stops listening when the signal aborts', async () => { - await wsServer.connected; - - const controller = new AbortController(); - const promise = asyncEvent.waitForAsyncData( - asyncPendingEvent, - controller.signal, - ); - const assertion = expect(promise).rejects.toMatchObject({ - name: 'AbortError', - }); - controller.abort(); - await assertion; - - // A late DONE event must not trigger a cached-data fetch: the listener - // was removed on abort, so no leak / stray request. - wsServer.send(JSON.stringify(asyncDoneEvent)); - await new Promise(resolve => { - setTimeout(resolve, 0); - }); - expect(fetchMock.callHistory.calls(CACHED_DATA_ENDPOINT)).toHaveLength(0); - }); - - test('rejects immediately if the signal is already aborted', async () => { - await wsServer.connected; - - const controller = new AbortController(); - controller.abort(); - - await expect( - asyncEvent.waitForAsyncData(asyncPendingEvent, controller.signal), - ).rejects.toMatchObject({ name: 'AbortError' }); - }); - - test('forwards the abort signal to the cached-data download', async () => { - await wsServer.connected; - - const getSpy = jest.spyOn(SupersetClient, 'get'); - const controller = new AbortController(); - - const promise = asyncEvent.waitForAsyncData( - asyncPendingEvent, - controller.signal, - ); - wsServer.send(JSON.stringify(asyncDoneEvent)); - await expect(promise).resolves.toEqual([chartData]); - - // The cached-result download must receive the signal so it can be - // cancelled if the caller aborts mid-fetch. - expect(getSpy).toHaveBeenCalledWith( - expect.objectContaining({ signal: controller.signal }), - ); - getSpy.mockRestore(); - }); +// Queue of status_changes responses the polling loop drains in order. The first +// is the baseline (empty), then each poll consumes the next. +let statusResponses: { statuses: Record; cursor: string }[]; + +const queueStatuses = ( + ...batches: Record[] +): void => { + statusResponses = [ + { statuses: {}, cursor: '2020-01-01T00:00:00' }, // baseline + ...batches.map((statuses, i) => ({ + statuses, + cursor: `2020-01-01T00:00:0${i + 1}`, + })), + ]; +}; + +beforeEach(() => { + mockedIsFeatureEnabled.mockImplementation( + featureFlag => featureFlag === 'GLOBAL_ASYNC_QUERIES', + ); + fetchMock.get(STATUS_CHANGES_ENDPOINT, () => { + const next = statusResponses.shift(); + return { status: 200, body: next ?? { statuses: {}, cursor: null } }; }); + fetchMock.post(CANCEL_ENDPOINT, { status: 200, body: { action: 'aborted' } }); +}); + +afterEach(() => { + fetchMock.clearHistory().removeRoutes(); + mockedIsFeatureEnabled.mockRestore(); +}); + +test('re-issues the request once every query task succeeds', async () => { + queueStatuses({ + 'task-1': { status: 'success' }, + 'task-2': { status: 'success' }, + }); + asyncEvent.init(config); + + const refetch = jest.fn().mockResolvedValue([{ rows: 1 }]); + const result = await asyncEvent.waitForAsyncData( + { task_ids: ['task-1', 'task-2'] }, + refetch, + ); + + expect(refetch).toHaveBeenCalledTimes(1); + expect(result).toEqual([{ rows: 1 }]); +}); + +test('waits for every task before re-issuing', async () => { + // task-2 only succeeds in the second poll batch. + queueStatuses( + { 'task-1': { status: 'success' } }, + { 'task-2': { status: 'success' } }, + ); + asyncEvent.init(config); + + const refetch = jest.fn().mockResolvedValue([{ rows: 2 }]); + await asyncEvent.waitForAsyncData( + { task_ids: ['task-1', 'task-2'] }, + refetch, + ); + + expect(refetch).toHaveBeenCalledTimes(1); +}); + +test('rejects and does not re-issue when a task fails', async () => { + queueStatuses({ + 'task-1': { status: 'success' }, + 'task-2': { status: 'failure' }, + }); + asyncEvent.init(config); + + const refetch = jest.fn(); + await expect( + asyncEvent.waitForAsyncData({ task_ids: ['task-1', 'task-2'] }, refetch), + ).rejects.toThrow(); + expect(refetch).not.toHaveBeenCalled(); +}); + +test('resolves immediately for an empty task list', async () => { + queueStatuses(); + asyncEvent.init(config); + + const refetch = jest.fn().mockResolvedValue([]); + await asyncEvent.waitForAsyncData({ task_ids: [] }, refetch); + expect(refetch).toHaveBeenCalledTimes(1); +}); + +test('aborting cancels the tasks and rejects with AbortError', async () => { + queueStatuses(); // tasks never resolve on their own + asyncEvent.init(config); + + const controller = new AbortController(); + const refetch = jest.fn(); + const promise = asyncEvent.waitForAsyncData( + { task_ids: ['task-1'] }, + refetch, + controller.signal, + ); + controller.abort(); + + await expect(promise).rejects.toThrow('Aborted'); + expect(refetch).not.toHaveBeenCalled(); + expect(fetchMock.callHistory.calls(CANCEL_ENDPOINT)).toHaveLength(1); +}); + +test('settles every request awaiting a deduplicated shared task', async () => { + // Two concurrent chart requests share the same (deduplicated) task uuid; both + // must resolve when it completes — the later waiter must not overwrite the first. + queueStatuses({ shared: { status: 'success' } }); + asyncEvent.init(config); + + const refetchA = jest.fn().mockResolvedValue([{ chart: 'a' }]); + const refetchB = jest.fn().mockResolvedValue([{ chart: 'b' }]); + const [a, b] = await Promise.all([ + asyncEvent.waitForAsyncData({ task_ids: ['shared'] }, refetchA), + asyncEvent.waitForAsyncData({ task_ids: ['shared'] }, refetchB), + ]); + + expect(refetchA).toHaveBeenCalledTimes(1); + expect(refetchB).toHaveBeenCalledTimes(1); + expect(a).toEqual([{ chart: 'a' }]); + expect(b).toEqual([{ chart: 'b' }]); }); diff --git a/superset-frontend/src/middleware/asyncEvent.ts b/superset-frontend/src/middleware/asyncEvent.ts index d82d613ac3a..c64ec4ee622 100644 --- a/superset-frontend/src/middleware/asyncEvent.ts +++ b/superset-frontend/src/middleware/asyncEvent.ts @@ -17,298 +17,204 @@ * under the License. */ import { - ensureIsArray, isFeatureEnabled, FeatureFlag, makeApi, SupersetClient, - getClientErrorObject, - parseErrorJson, - SupersetError, } from '@superset-ui/core'; import { logging } from '@apache-superset/core/utils'; import getBootstrapData from 'src/utils/getBootstrapData'; -type AsyncEvent = { - id?: string | null; - channel_id: string; - job_id: string; - user_id?: string; - status: string; - errors?: SupersetError[]; - result_url: string | null; +// The GTF task type chart-data queries run under (see +// superset/tasks/async_queries.py CHART_QUERY_TASK). Polling is filtered to this +// type so a dashboard only tracks its own chart-data work, not every task. +const CHART_QUERY_TASK_TYPE = 'superset.query_object_v1'; +const STATUS_CHANGES_URL = '/api/v1/task/status_changes'; + +// Terminal GTF task statuses (mirror superset_core.tasks.types.TaskStatus). +const STATUS_SUCCESS = 'success'; +const TERMINAL_STATUSES = new Set([ + STATUS_SUCCESS, + 'failure', + 'aborted', + 'timed_out', +]); + +type TaskStatusChange = { status: string; progress: number | null }; +type StatusChangesResponse = { + statuses: Record; + cursor: string | null; }; -type CachedDataResponse = { - status: string; - data: any; -}; +// The 202 body from POST /chart/data when async: the query tasks to await. +export type AsyncJob = { task_ids: string[] }; + type AppConfig = Record; -type ListenerFn = (asyncEvent: AsyncEvent) => Promise; -const TRANSPORT_POLLING = 'polling'; -const TRANSPORT_WS = 'ws'; -const JOB_STATUS = { - PENDING: 'pending', - RUNNING: 'running', - ERROR: 'error', - DONE: 'done', +type Waiter = { + taskIds: string[]; + pending: Set; + failed: boolean; + // Re-issue the original chart-data request once every task has succeeded; the + // per-query DATA cache is now warm, so it returns synchronously (200). + resolve: () => void; + reject: (error: unknown) => void; + signal?: AbortSignal; + onAbort?: () => void; }; -const LOCALSTORAGE_KEY = 'last_async_event_id'; -const POLLING_URL = '/api/v1/async_event/'; -const MAX_RETRIES = 6; -const RETRY_DELAY = 100; -// Cap for the exponential backoff applied when polling requests fail -// repeatedly (e.g. expired session, server or network errors) -const MAX_ERROR_POLLING_DELAY_MS = 60000; let config: AppConfig; -let transport: string; let pollingDelayMs: number; let pollingTimeoutId: number; -let listenersByJobId: Map; -let retriesByJobId: Map; -let lastReceivedEventId: string | null | undefined; -let consecutivePollingErrorCount = 0; -// Incremented on every init() so polling invocations that are already -// awaiting a fetch when re-init happens can detect they are stale and -// stop, instead of mutating fresh state or scheduling a second loop +// Registry of in-flight waiters keyed by every task uuid they await, so a single +// shared poll loop fans status changes out to whichever requests are awaiting them. +// A SHARED task can be deduplicated across concurrent chart requests, so each task +// id maps to a *set* of waiters (never overwrite an earlier subscriber). +let waitersByTaskId: Map>; +// Server-issued watermark: fetched as a baseline at init (before any chart query +// is triggered) so no task created afterwards is missed, then advanced by each +// poll. Always the server's own clock, never the browser's. +let cursor: string | null; +let baselineReady: Promise | null; +// Incremented on every init() so an in-flight poll can detect it is stale and +// stop instead of scheduling a second loop or mutating fresh state. let pollingGeneration = 0; -const addListener = (id: string, fn: ListenerFn) => { - listenersByJobId.set(id, fn); -}; - -const removeListener = (id: string) => { - if (!listenersByJobId.has(id)) return; - listenersByJobId.delete(id); -}; - -const fetchCachedData = async ( - asyncEvent: AsyncEvent, - signal?: AbortSignal, -): Promise => { - let status = 'success'; - let data; - try { - const { json } = await SupersetClient.get({ - endpoint: String(asyncEvent.result_url), - signal, - }); - data = 'result' in json ? json.result : json; - } catch (response) { - status = 'error'; - data = await getClientErrorObject(response); - } - - return { status, data }; -}; - -const cancelAsyncJob = (jobId: string) => { - // Best-effort server-side cancel; the request stops the running Celery task - // so it no longer consumes warehouse resources. Failures are non-fatal: the - // client has already stopped waiting on the job. - SupersetClient.post({ - endpoint: `/api/v1/async_event/${jobId}/cancel`, - }).catch(error => { - logging.warn('Failed to cancel async job', jobId, error); - }); -}; - -export const waitForAsyncData = async ( - asyncResponse: AsyncEvent, - signal?: AbortSignal, -) => - new Promise((resolve, reject) => { - const jobId = asyncResponse.job_id; - - let onAbort: (() => void) | undefined; - const cleanup = () => { - removeListener(jobId); - if (onAbort && signal) { - signal.removeEventListener('abort', onAbort); - } - }; - - // Bail immediately if the caller has already aborted (e.g. the chart was - // unmounted before the job started), avoiding a leaked listener. - if (signal?.aborted) { - cancelAsyncJob(jobId); - reject(new DOMException('Aborted', 'AbortError')); - return; - } - - const listener = async (asyncEvent: AsyncEvent) => { - switch (asyncEvent.status) { - case JOB_STATUS.DONE: { - // Forward the signal so the cached-result download is cancelled too if - // the caller aborts mid-fetch, rather than wasting network/processing. - let { data, status } = await fetchCachedData(asyncEvent, signal); // eslint-disable-line prefer-const - data = ensureIsArray(data); - if (status === 'success') { - resolve(data); - } else { - reject(data); - } - // Terminal status: the promise is settled, so fully clean up. - cleanup(); - break; - } - case JOB_STATUS.ERROR: { - const err = parseErrorJson(asyncEvent); - reject(err); - // Terminal status: the promise is settled, so fully clean up. - cleanup(); - break; - } - default: { - // Non-terminal status (e.g., 'pending', 'running'): keep the listener - // registered so it can receive the eventual terminal event ('done', 'error'). - // Only cleanup happens on terminal states or abort. - logging.info( - 'received non-terminal event with status', - asyncEvent.status, - ); - } - } - }; - - // When the caller aborts (Stop pressed, chart superseded/unmounted), stop - // listening so the listener and its retained closure don't leak, and ask the - // server to cancel the job so it stops consuming warehouse resources. - if (signal) { - onAbort = () => { - cleanup(); - cancelAsyncJob(jobId); - reject(new DOMException('Aborted', 'AbortError')); - }; - signal.addEventListener('abort', onAbort, { once: true }); - } - - addListener(jobId, listener); - }); - -const fetchEvents = makeApi< - { last_id?: string | null }, - { result: AsyncEvent[] } +const fetchStatusChanges = makeApi< + { cursor?: string | null; task_type: string }, + StatusChangesResponse >({ method: 'GET', - endpoint: POLLING_URL, + endpoint: STATUS_CHANGES_URL, }); -const setLastId = (asyncEvent: AsyncEvent) => { - lastReceivedEventId = asyncEvent.id; - try { - localStorage.setItem(LOCALSTORAGE_KEY, lastReceivedEventId as string); - } catch (err) { - logging.warn('Error saving event Id to localStorage', err); - } -}; - -export const processEvents = async (events: AsyncEvent[]) => { - events.forEach((asyncEvent: AsyncEvent) => { - const jobId = asyncEvent.job_id; - const listener = listenersByJobId.get(jobId); - // `jobId` originates from server/WebSocket payloads, so the listener is - // resolved exclusively through a Map (never plain-object property access, - // which would expose the prototype chain), and we confirm the retrieved - // value is a registered function before dispatching the event to it. - if (typeof listener === 'function') { - listener(asyncEvent); - retriesByJobId.delete(jobId); - } else { - // handle race condition where event is received - // before listener is registered - const retries = (retriesByJobId.get(jobId) ?? 0) + 1; - retriesByJobId.set(jobId, retries); - - if (retries <= MAX_RETRIES) { - setTimeout(() => { - processEvents([asyncEvent]); - }, RETRY_DELAY * retries); - } else { - retriesByJobId.delete(jobId); - logging.warn('listener not found for job_id', asyncEvent.job_id); - } - } - setLastId(asyncEvent); +const cancelTask = (taskId: string) => { + // Best-effort server-side cancel so an abandoned query stops consuming + // warehouse resources. Failures are non-fatal: the client has already stopped + // waiting on the task. + SupersetClient.post({ + endpoint: `/api/v1/task/${taskId}/cancel`, + }).catch(error => { + logging.warn('Failed to cancel task', taskId, error); }); }; -const getPollingDelay = () => { - if (!consecutivePollingErrorCount) return pollingDelayMs; - const backoffDelayMs = pollingDelayMs * 2 ** consecutivePollingErrorCount; - return Math.max( +// Drop a waiter from the registry entry of every task it was awaiting, so a +// settled/aborted waiter never leaks and completion of one task can't re-touch it. +const unregister = (waiter: Waiter) => { + waiter.taskIds.forEach(taskId => { + const waiters = waitersByTaskId.get(taskId); + if (!waiters) return; + waiters.delete(waiter); + if (waiters.size === 0) waitersByTaskId.delete(taskId); + }); +}; + +const settle = (waiter: Waiter) => { + unregister(waiter); + if (waiter.signal && waiter.onAbort) { + waiter.signal.removeEventListener('abort', waiter.onAbort); + } + if (waiter.failed) { + waiter.reject( + new Error('One or more chart-data queries failed'), // surfaced via getClientErrorObject + ); + } else { + waiter.resolve(); + } +}; + +const applyStatus = (taskId: string, status: string) => { + const waiters = waitersByTaskId.get(taskId); + if (!waiters || !TERMINAL_STATUSES.has(status)) return; + // Settle every request awaiting this task, not just the most recent one. + [...waiters].forEach(waiter => { + waiter.pending.delete(taskId); + if (status !== STATUS_SUCCESS) waiter.failed = true; + if (waiter.pending.size === 0) settle(waiter); + }); + waitersByTaskId.delete(taskId); +}; + +const loadStatusChanges = async (generation: number) => { + if (generation !== pollingGeneration) return; + if (waitersByTaskId.size) { + try { + const { statuses, cursor: next } = await fetchStatusChanges({ + cursor, + task_type: CHART_QUERY_TASK_TYPE, + }); + if (generation !== pollingGeneration) return; + cursor = next; + Object.entries(statuses).forEach(([taskId, { status }]) => + applyStatus(taskId, status), + ); + } catch (err) { + if (generation !== pollingGeneration) return; + logging.warn(err); + } + } + // Reschedule from the tail so a slow request never overlaps the next tick. + pollingTimeoutId = window.setTimeout( + () => loadStatusChanges(generation), pollingDelayMs, - Math.min(backoffDelayMs, MAX_ERROR_POLLING_DELAY_MS), ); }; -const loadEventsFromApi = async () => { - const generation = pollingGeneration; - const eventArgs = lastReceivedEventId ? { last_id: lastReceivedEventId } : {}; - if (listenersByJobId.size) { - try { - const { result: events } = await fetchEvents(eventArgs); - if (generation !== pollingGeneration) return; - consecutivePollingErrorCount = 0; - if (events?.length) await processEvents(events); - } catch (err) { - if (generation !== pollingGeneration) return; - consecutivePollingErrorCount += 1; - logging.warn(err); +/** + * Await completion of an async chart-data job's query tasks, then re-issue the + * original request to read the now-cached results. + * + * Resolves with the fresh `QueryData[]` once every task has succeeded (the + * caller's `refetch` returns synchronously from the warm per-query cache); + * rejects if any task ends in a non-success terminal state, or with an + * AbortError if the caller aborts (which also cancels the outstanding tasks). + */ +export const waitForAsyncData = async ( + asyncJob: AsyncJob, + refetch: () => Promise, + signal?: AbortSignal, +): Promise => { + const taskIds = asyncJob.task_ids ?? []; + if (baselineReady) await baselineReady; + + await new Promise((resolve, reject) => { + if (signal?.aborted) { + taskIds.forEach(cancelTask); + reject(new DOMException('Aborted', 'AbortError')); + return; } - } - - if (generation !== pollingGeneration) return; - if (transport === TRANSPORT_POLLING) { - pollingTimeoutId = window.setTimeout(loadEventsFromApi, getPollingDelay()); - } -}; - -const wsConnectMaxRetries = 6; -const wsConnectErrorDelay = 2500; -let wsConnectRetries = 0; -let wsConnectTimeout: any; -let ws: WebSocket; - -const wsConnect = (): void => { - let url = config.GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL; - if (lastReceivedEventId) url += `?last_id=${lastReceivedEventId}`; - ws = new WebSocket(url); - - ws.addEventListener('open', () => { - logging.log('WebSocket connected'); - clearTimeout(wsConnectTimeout); - wsConnectRetries = 0; - }); - - ws.addEventListener('close', () => { - wsConnectTimeout = setTimeout(() => { - wsConnectRetries += 1; - if (wsConnectRetries <= wsConnectMaxRetries) { - wsConnect(); - } else { - logging.warn('WebSocket not available, falling back to async polling'); - loadEventsFromApi(); + const waiter: Waiter = { + taskIds, + pending: new Set(taskIds), + failed: false, + resolve, + reject, + signal, + }; + if (signal) { + waiter.onAbort = () => { + unregister(waiter); + taskIds.forEach(cancelTask); + reject(new DOMException('Aborted', 'AbortError')); + }; + signal.addEventListener('abort', waiter.onAbort, { once: true }); + } + if (!taskIds.length) { + settle(waiter); + return; + } + taskIds.forEach(taskId => { + let waiters = waitersByTaskId.get(taskId); + if (!waiters) { + waiters = new Set(); + waitersByTaskId.set(taskId, waiters); } - }, wsConnectErrorDelay); + waiters.add(waiter); + }); }); - ws.addEventListener('error', () => { - // https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState - if (ws.readyState < 2) ws.close(); - }); - - ws.addEventListener('message', async event => { - let events: AsyncEvent[] = []; - try { - events = [JSON.parse(event.data)]; - await processEvents(events); - } catch (err) { - logging.warn(err); - } - }); + return refetch(); }; export const init = (appConfig?: AppConfig) => { @@ -316,27 +222,28 @@ export const init = (appConfig?: AppConfig) => { if (pollingTimeoutId) clearTimeout(pollingTimeoutId); if (!isFeatureEnabled(FeatureFlag.GlobalAsyncQueries)) return; - listenersByJobId = new Map(); - retriesByJobId = new Map(); - lastReceivedEventId = null; - consecutivePollingErrorCount = 0; + const generation = pollingGeneration; + waitersByTaskId = new Map(); + cursor = null; config = appConfig || getBootstrapData().common.conf; - transport = config.GLOBAL_ASYNC_QUERIES_TRANSPORT || TRANSPORT_POLLING; pollingDelayMs = config.GLOBAL_ASYNC_QUERIES_POLLING_DELAY || 500; - try { - lastReceivedEventId = localStorage.getItem(LOCALSTORAGE_KEY); - } catch (err) { - logging.warn('Failed to fetch last event Id from localStorage'); - } - - if (transport === TRANSPORT_POLLING) { - loadEventsFromApi(); - } - if (transport === TRANSPORT_WS) { - wsConnect(); - } + // Establish a baseline cursor before any chart query is triggered, so tasks + // created afterwards are all caught by the changed-since poll, then start the + // shared poll loop from that watermark. + baselineReady = fetchStatusChanges({ task_type: CHART_QUERY_TASK_TYPE }) + .then(({ cursor: baseline }) => { + if (generation === pollingGeneration) cursor = baseline; + }) + .catch(err => { + // A missing baseline just means the first poll starts from "everything + // changed so far"; the >= cursor semantics still catch our tasks. + logging.warn('Failed to fetch async baseline cursor', err); + }) + .finally(() => { + loadStatusChanges(generation); + }); }; init(); diff --git a/superset/async_events/api.py b/superset/async_events/api.py deleted file mode 100644 index 7968f956b54..00000000000 --- a/superset/async_events/api.py +++ /dev/null @@ -1,179 +0,0 @@ -# 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 logging -import uuid - -from flask import request, Response -from flask_appbuilder import expose -from flask_appbuilder.api import safe -from flask_appbuilder.security.decorators import permission_name, protect - -from superset.async_events.async_query_manager import ( - AsyncQueryJobException, - AsyncQueryTokenException, -) -from superset.extensions import async_query_manager, event_logger -from superset.utils.core import get_user_id -from superset.views.base_api import BaseSupersetApi, statsd_metrics - -logger = logging.getLogger(__name__) - - -class AsyncEventsRestApi(BaseSupersetApi): - resource_name = "async_event" - allow_browser_login = True - - @expose("/", methods=("GET",)) - @event_logger.log_this - @protect() - @safe - @statsd_metrics - @permission_name("list") - def events(self) -> Response: - """ - Read off of the Redis async events stream, using the user's JWT token and - optional query params for last event received. - --- - get: - summary: Read off of the Redis events stream - description: >- - Reads off of the Redis events stream, using the user's JWT token and - optional query params for last event received. - parameters: - - in: query - name: last_id - description: Last ID received by the client - schema: - type: string - responses: - 200: - description: Async event results - content: - application/json: - schema: - type: object - properties: - result: - type: array - items: - type: object - properties: - id: - type: string - channel_id: - type: string - job_id: - type: string - user_id: - type: integer - status: - type: string - errors: - type: array - items: - type: object - result_url: - type: string - 401: - $ref: '#/components/responses/401' - 500: - $ref: '#/components/responses/500' - """ - try: - async_channel_id = async_query_manager.parse_channel_id_from_request( - request - ) - last_event_id = request.args.get("last_id") - events = async_query_manager.read_events(async_channel_id, last_event_id) - - except AsyncQueryTokenException: - return self.response_401() - - return self.response(200, result=events) - - @expose("//cancel", methods=("POST",)) - @event_logger.log_this - @protect() - @safe - @statsd_metrics - @permission_name("cancel") - def cancel(self, job_id: str) -> Response: - """Cancel a running async query job. - --- - post: - summary: Cancel a running async query job - description: >- - Revokes the Celery task backing an in-flight async query. The - caller is authorized against the job's original owner (channel and - user), both resolved server-side from the request, so a client - cannot cancel a job it did not submit. - parameters: - - in: path - name: job_id - required: true - description: The job ID returned when the async query was submitted - schema: - type: string - responses: - 200: - description: Job cancelled - content: - application/json: - schema: - type: object - properties: - result: - type: object - properties: - job_id: - type: string - status: - type: string - 400: - $ref: '#/components/responses/400' - 401: - $ref: '#/components/responses/401' - 403: - $ref: '#/components/responses/403' - 404: - $ref: '#/components/responses/404' - 500: - $ref: '#/components/responses/500' - """ - try: - uuid.UUID(job_id) - except ValueError: - return self.response_400(message="Invalid job ID") - - try: - async_channel_id = async_query_manager.parse_channel_id_from_request( - request - ) - except AsyncQueryTokenException: - return self.response_401() - - try: - async_query_manager.cancel_job(job_id, async_channel_id, get_user_id()) - except AsyncQueryTokenException: - return self.response_403() - except AsyncQueryJobException: - return self.response_404() - - return self.response( - 200, - result={"job_id": job_id, "status": async_query_manager.STATUS_CANCELLED}, - ) diff --git a/superset/async_events/async_query_manager.py b/superset/async_events/async_query_manager.py deleted file mode 100644 index bdf94c4fb45..00000000000 --- a/superset/async_events/async_query_manager.py +++ /dev/null @@ -1,533 +0,0 @@ -# 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. -from __future__ import annotations - -import hashlib -import hmac -import logging -import uuid -from datetime import datetime, timedelta, timezone -from typing import Any, Literal, Optional, TYPE_CHECKING - -import jwt -from flask import Flask, Request, request, Response, session - -from superset.async_events.cache_backend import ( - RedisCacheBackend, - RedisSentinelCacheBackend, -) -from superset.coordination.base import CoordinationService -from superset.utils import json -from superset.utils.core import get_user_id - -if TYPE_CHECKING: - from superset.security.guest_token import GuestUser - -logger = logging.getLogger(__name__) - - -class CacheBackendNotInitialized(Exception): # noqa: N818 - pass - - -class AsyncQueryTokenException(Exception): # noqa: N818 - pass - - -class UnsupportedCacheBackendError(Exception): # noqa: N818 - pass - - -class AsyncQueryJobException(Exception): # noqa: N818 - pass - - -def build_job_metadata( - channel_id: str, job_id: str, user_id: Optional[int], **kwargs: Any -) -> dict[str, Any]: - return { - "channel_id": channel_id, - "job_id": job_id, - "user_id": user_id, - "status": kwargs.get("status"), - "errors": kwargs.get("errors", []), - "result_url": kwargs.get("result_url"), - } - - -def parse_event(event_data: tuple[str, dict[str, Any]]) -> dict[str, Any]: - event_id = event_data[0] - event_payload = event_data[1]["data"] - return {"id": event_id, **json.loads(event_payload)} - - -def increment_id(entry_id: str) -> str: - # redis stream IDs are in this format: '1607477697866-0' - try: - prefix, last = entry_id[:-1], int(entry_id[-1]) - return prefix + str(last + 1) - except Exception: # pylint: disable=broad-except - return entry_id - - -def get_cache_backend( - config: dict[str, Any], -) -> RedisCacheBackend | RedisSentinelCacheBackend: - """Build a coordination backend from the deprecated GAQ cache config. - - DEPRECATED: retained only so Global Async Queries can run on its own dedicated - ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` during the deprecation window. Removed in - Superset 8.0, when GAQ moves onto ``DISTRIBUTED_COORDINATION_CONFIG``. - """ - cache_config = config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}) - cache_type = cache_config.get("CACHE_TYPE") - - if cache_type == "RedisCache": - return RedisCacheBackend.from_config(cache_config) - - if cache_type == "RedisSentinelCache": - return RedisSentinelCacheBackend.from_config(cache_config) - - raise UnsupportedCacheBackendError("Unsupported cache backend configuration") - - -class AsyncQueryManager: - MAX_EVENT_COUNT = 100 - STATUS_PENDING = "pending" - STATUS_RUNNING = "running" - STATUS_ERROR = "error" - STATUS_DONE = "done" - STATUS_CANCELLED = "cancelled" - # Redis key prefix (within the GAQ stream namespace) for the per-job record - # that authorizes cancellation and flags a job as cancelled for the worker. - _JOB_REGISTRY_PREFIX = "job-cancel:" - # Emit the dedicated-backend deprecation notice at most once per process. - _legacy_backend_warning_emitted: bool = False - - def __init__(self) -> None: - super().__init__() - self._stream_prefix: str = "" - self._stream_limit: Optional[int] - self._stream_limit_firehose: Optional[int] - # Global Async Queries owns its coordination backend separately from the - # shared coordinator (see init_app); resolved there and passed explicitly - # to CoordinationService's primitives. - self._gaq_backend: RedisCacheBackend | RedisSentinelCacheBackend | None = None - self._jwt_cookie_name: str = "" - self._jwt_cookie_secure: bool = False - self._jwt_cookie_domain: Optional[str] - self._jwt_cookie_samesite: Optional[Literal["None", "Lax", "Strict"]] = None - self._jwt_secret: str - self._jwt_expiration_seconds: int = 0 - self._load_chart_data_into_cache_job: Any = None - # pylint: disable=invalid-name - - def init_app(self, app: Flask) -> None: - cache_type = app.config.get("CACHE_CONFIG", {}).get("CACHE_TYPE") - data_cache_type = app.config.get("DATA_CACHE_CONFIG", {}).get("CACHE_TYPE") - if cache_type in [None, "null"] or data_cache_type in [None, "null"]: - raise Exception( # pylint: disable=broad-exception-raised - """ - Cache backends (CACHE_CONFIG, DATA_CACHE_CONFIG) must be configured - and non-null in order to enable async queries - """ - ) - - if not ( - app.config.get("DISTRIBUTED_COORDINATION_CONFIG") - or app.config.get("GLOBAL_ASYNC_QUERIES_CACHE_BACKEND", {}).get( - "CACHE_TYPE" - ) - ): - raise UnsupportedCacheBackendError( - "Global async queries require a coordination backend; configure " - "DISTRIBUTED_COORDINATION_CONFIG (GLOBAL_ASYNC_QUERIES_CACHE_BACKEND " - "is deprecated)." - ) - - # Global Async Queries share the coordinator's backend whenever - # DISTRIBUTED_COORDINATION_CONFIG is configured, so an 8.0-leaning deployment - # keeps a single coordination config rather than maintaining a separate one. - # Only when no coordinator is configured does GAQ fall back to its dedicated - # (deprecated) GLOBAL_ASYNC_QUERIES_CACHE_BACKEND, emitting a one-time - # deprecation warning. That dedicated backend is removed in Superset 8.0, when - # DISTRIBUTED_COORDINATION_CONFIG becomes the only option. Either way GAQ passes - # its resolved backend explicitly to CoordinationService's primitives, keeping - # its stream/pub-sub traffic scoped to the connection it resolved here. - if (coordinator := CoordinationService.get_backend()) is not None: - self._gaq_backend = coordinator - else: - if not AsyncQueryManager._legacy_backend_warning_emitted: - logger.warning( - "Global Async Queries is running on the deprecated " - "GLOBAL_ASYNC_QUERIES_CACHE_BACKEND because " - "DISTRIBUTED_COORDINATION_CONFIG is not configured. Configure " - "DISTRIBUTED_COORDINATION_CONFIG to consolidate coordination " - "(distributed locks, task framework, pub/sub, and the GAQ event " - "streams) onto a single connection; the dedicated backend is " - "removed in Superset 8.0." - ) - AsyncQueryManager._legacy_backend_warning_emitted = True - self._gaq_backend = get_cache_backend(app.config) - - if len(app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"]) < 32: - raise AsyncQueryTokenException( - "Please provide a JWT secret at least 32 bytes long" - ) - - self._stream_prefix = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] - self._stream_limit = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT"] - self._stream_limit_firehose = app.config[ - "GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT_FIREHOSE" - ] - self._jwt_cookie_name = app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"] - self._jwt_cookie_secure = app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SECURE"] - self._jwt_cookie_samesite = app.config[ - "GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SAMESITE" - ] - self._jwt_cookie_domain = app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_DOMAIN"] - self._jwt_secret = app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"] - self._jwt_expiration_seconds = app.config[ - "GLOBAL_ASYNC_QUERIES_JWT_EXPIRATION_SECONDS" - ] - - if app.config["GLOBAL_ASYNC_QUERIES_REGISTER_REQUEST_HANDLERS"]: - self.register_request_handlers(app) - - # pylint: disable=import-outside-toplevel - from superset.tasks.async_queries import load_chart_data_into_cache - - self._load_chart_data_into_cache_job = load_chart_data_into_cache - - def register_request_handlers(self, app: Flask) -> None: - @app.after_request - def validate_session(response: Response) -> Response: - # pylint: disable=import-outside-toplevel - from superset import security_manager - - # Guest users (embedded dashboards) are typically loaded from a - # third-party context where session cookies are unreliable, so the - # async channel is derived deterministically from the guest token - # in `parse_channel_id_from_request` and the JWT cookie is not - # required. - if security_manager.get_current_guest_user_if_guest(): - return response - - user_id = get_user_id() - - reset_token = ( - not request.cookies.get(self._jwt_cookie_name) - or "async_channel_id" not in session - or "async_user_id" not in session - or user_id != session["async_user_id"] - ) - - if reset_token: - async_channel_id = str(uuid.uuid4()) - session["async_channel_id"] = async_channel_id - session["async_user_id"] = user_id - - # Conditionally include 'sub' claim only when user_id is present. - # RFC 7519 specifies 'sub' as optional; when present it must be - # a string, so omit it entirely for guest/anonymous users. - now = datetime.now(tz=timezone.utc) - payload = { - "channel": async_channel_id, - "exp": now + timedelta(seconds=self._jwt_expiration_seconds), - } - if user_id is not None: - payload["sub"] = str(user_id) - token = jwt.encode( - payload, - self._jwt_secret, - algorithm="HS256", - ) - - response.set_cookie( - self._jwt_cookie_name, - value=token, - httponly=True, - secure=self._jwt_cookie_secure, - domain=self._jwt_cookie_domain, - samesite=self._jwt_cookie_samesite, - max_age=self._jwt_expiration_seconds, - ) - - return response - - def get_guest_user_channel_id(self, guest_user: GuestUser) -> str: - """ - Derive a deterministic async channel ID for a guest user. - - Embedded guest sessions cannot reliably rely on the async-token cookie - because cross-origin cookies are blocked or stripped by modern browsers - when running inside a third-party iframe. Using an HMAC over stable - guest-token claims yields a per-token channel that the chart data - request, the celery worker, and the polling endpoint can all derive - without needing a cookie. The HMAC is keyed with the configured JWT - secret so the value is unguessable to outside callers. - """ - token = guest_user.guest_token - # ``iat`` uniquely identifies a guest token issuance, so it provides - # per-token isolation while remaining stable across the lifetime of a - # single embedded session. - message = json.dumps( - { - "user": token.get("user"), - "resources": token.get("resources"), - "iat": token.get("iat"), - "exp": token.get("exp"), - "aud": token.get("aud"), - # ``datasets`` and ``rev`` are optional scope claims, so tokens - # that differ only in their dataset allowlist or revocation - # version still derive distinct channels. - "datasets": token.get("datasets"), - "rev": token.get("rev"), - }, - sort_keys=True, - ).encode("utf-8") - digest = hmac.new( - self._jwt_secret.encode("utf-8"), message, hashlib.sha256 - ).hexdigest() - return f"guest-{digest}" - - def parse_channel_id_from_request(self, req: Request) -> str: - # pylint: disable=import-outside-toplevel - from superset import security_manager - - if guest_user := security_manager.get_current_guest_user_if_guest(): - return self.get_guest_user_channel_id(guest_user) - - token = req.cookies.get(self._jwt_cookie_name) - if not token: - raise AsyncQueryTokenException("Token not preset") - - try: - return jwt.decode(token, self._jwt_secret, algorithms=["HS256"])["channel"] - except Exception as ex: - logger.warning("Parse jwt failed", exc_info=True) - raise AsyncQueryTokenException("Failed to parse token") from ex - - def init_job(self, channel_id: str, user_id: Optional[int]) -> dict[str, Any]: - job_id = str(uuid.uuid4()) - self._register_cancellable_job(job_id, channel_id, user_id) - return build_job_metadata( - channel_id, job_id, user_id, status=self.STATUS_PENDING - ) - - def _job_registry_key(self, job_id: str) -> str: - return f"{self._stream_prefix}{self._JOB_REGISTRY_PREFIX}{job_id}" - - def _register_cancellable_job( - self, job_id: str, channel_id: str, user_id: Optional[int] - ) -> None: - """ - Persist the identity a later cancel request must match. Keyed by - ``job_id`` (also the Celery task id — see ``submit_chart_data_job``) so - the cancel endpoint can authorize the caller against the job's original - owner without trusting the client-supplied id. Expires with the JWT so - it never outlives the job it guards. - """ - # The cancel registry is an optimization: skip it when GAQ has no - # coordination backend configured. When a backend is configured, a write - # failure is not swallowed here — it surfaces to the caller and fails job - # submission, matching the behavior of the surrounding stream writes. - if self._gaq_backend is None: - return - CoordinationService.set_value( - self._job_registry_key(job_id), - json.dumps({"channel_id": channel_id, "user_id": user_id}), - ttl=self._jwt_expiration_seconds or None, - backend=self._gaq_backend, - ) - - def submit_chart_data_job( - self, - channel_id: str, - form_data: dict[str, Any], - user_id: Optional[int] = None, - ) -> dict[str, Any]: - # pylint: disable=import-outside-toplevel - from superset import security_manager - - # if it's guest user, we want to pass the guest token to the celery task - # chart data cache key is calculated based on the current user - # this way we can keep the cache key consistent between sync and async command - # so that it can be looked up consistently - job_metadata = self.init_job(channel_id, user_id) - self._load_chart_data_into_cache_job.apply_async( - args=[ - {**job_metadata, "guest_token": guest_user.guest_token} - if (guest_user := security_manager.get_current_guest_user_if_guest()) - else job_metadata, - form_data, - ], - # Use job_id as the Celery task id so the cancel endpoint can revoke - # the running task by the id the client already holds. - task_id=job_metadata["job_id"], - expires=self._jwt_expiration_seconds, - ) - return job_metadata - - def read_events( - self, channel: str, last_id: Optional[str] - ) -> list[Optional[dict[str, Any]]]: - if self._gaq_backend is None: - raise CacheBackendNotInitialized("Cache backend not initialized") - - stream_name = f"{self._stream_prefix}{channel}" - start_id = increment_id(last_id) if last_id else "-" - results = CoordinationService.stream_range( - stream_name, start_id, "+", self.MAX_EVENT_COUNT, backend=self._gaq_backend - ) - # Decode bytes to strings: the coordination Redis backends do not enable - # decode_responses, so stream_range returns raw bytes. - decoded_results = [ - ( - event_id.decode("utf-8"), - { - key.decode("utf-8"): value.decode("utf-8") - for key, value in event_data.items() - }, - ) - for event_id, event_data in results - ] - return [] if not decoded_results else list(map(parse_event, decoded_results)) - - def update_job( - self, job_metadata: dict[str, Any], status: str, **kwargs: Any - ) -> None: - if self._gaq_backend is None: - raise CacheBackendNotInitialized("Cache backend not initialized") - - if "channel_id" not in job_metadata: - raise AsyncQueryJobException("No channel ID specified") - - if "job_id" not in job_metadata: - raise AsyncQueryJobException("No job ID specified") - - updates = {"status": status, **kwargs} - event_data = {"data": json.dumps({**job_metadata, **updates})} - - full_stream_name = f"{self._stream_prefix}full" - scoped_stream_name = f"{self._stream_prefix}{job_metadata['channel_id']}" - - logger.debug("********** logging event data to stream %s", scoped_stream_name) - logger.debug(event_data) - - # Drop the cancel record before announcing the result, so a cancel that - # arrives once the job is done finds nothing to flag and is reported as - # such instead of contradicting the event below. A cancelled job keeps - # its record until the TTL: the worker still has to recognize the - # SIGUSR1 it is about to receive as a cancellation. - if status in (self.STATUS_DONE, self.STATUS_ERROR): - if job_id := job_metadata.get("job_id"): - CoordinationService.delete_value( - self._job_registry_key(job_id), backend=self._gaq_backend - ) - - CoordinationService.stream_add( - scoped_stream_name, - event_data, - "*", - self._stream_limit, - backend=self._gaq_backend, - ) - CoordinationService.stream_add( - full_stream_name, - event_data, - "*", - self._stream_limit_firehose, - backend=self._gaq_backend, - ) - - def is_job_cancelled(self, job_id: str) -> bool: - """ - Whether ``cancel_job`` has flagged this job for cancellation. - - Called from the worker's exception handler, so any cache failure is - swallowed and treated as "not cancelled" — a Redis blip must never mask - the original error (e.g. a genuine timeout) with a connection error. - """ - if self._gaq_backend is None: - return False - try: - raw = CoordinationService.get_value( - self._job_registry_key(job_id), backend=self._gaq_backend - ) - if raw is None: - return False - return bool(json.loads(raw).get("cancelled")) - except Exception: # pylint: disable=broad-except - logger.warning( - "Failed to read cancellation flag for job %s", job_id, exc_info=True - ) - return False - - def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int]) -> None: - """ - Authorize and cancel a running async job. - - The caller's ``channel_id`` and ``user_id`` (resolved server-side from - the request, never taken from the client) must match the job's original - owner. The terminal ``STATUS_CANCELLED`` event is emitted here rather - than by the worker, which never runs for a task revoked while it was - still queued; the worker only logs the cancellation it is told about - through the flag, so a job still gets exactly one terminal event. - - :raises AsyncQueryJobException: the job is unknown or already terminal - :raises AsyncQueryTokenException: the caller does not own the job - """ - if self._gaq_backend is None: - raise CacheBackendNotInitialized("Cache backend not initialized") - - key = self._job_registry_key(job_id) - raw = CoordinationService.get_value(key, backend=self._gaq_backend) - if raw is None: - raise AsyncQueryJobException("Job not found or already completed") - - record = json.loads(raw) - if record.get("channel_id") != channel_id or record.get("user_id") != user_id: - raise AsyncQueryTokenException("Not authorized to cancel this job") - - # Flag before revoking so the worker's timeout handler, which may fire - # almost immediately, reliably sees the cancellation. Write only if the - # key still exists (``xx``): if the job finished and cleared its record - # between the read above and here, don't recreate a stale record or - # revoke a task that is already gone — report it as not found instead. - flagged = CoordinationService.set_value( - key, - json.dumps({**record, "cancelled": True}), - ttl=self._jwt_expiration_seconds or None, - if_present=True, - backend=self._gaq_backend, - ) - if not flagged: - raise AsyncQueryJobException("Job not found or already completed") - - # pylint: disable=import-outside-toplevel - from superset.extensions import celery_app - - # SIGUSR1 raises SoftTimeLimitExceeded inside the running task rather - # than hard-killing the process, so it unwinds through the task's - # exception handling instead of dying mid-query. - celery_app.control.revoke(job_id, terminate=True, signal="SIGUSR1") - - self.update_job( - build_job_metadata(channel_id, job_id, user_id), - self.STATUS_CANCELLED, - ) diff --git a/superset/async_events/async_query_manager_factory.py b/superset/async_events/async_query_manager_factory.py deleted file mode 100644 index 2e05f386034..00000000000 --- a/superset/async_events/async_query_manager_factory.py +++ /dev/null @@ -1,35 +0,0 @@ -# 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. - -from flask import Flask - -from superset.async_events.async_query_manager import AsyncQueryManager -from superset.utils.class_utils import load_class_from_name - - -class AsyncQueryManagerFactory: - def __init__(self) -> None: - self._async_query_manager: AsyncQueryManager = None # type: ignore - - def init_app(self, app: Flask) -> None: - self._async_query_manager = load_class_from_name( - app.config["GLOBAL_ASYNC_QUERY_MANAGER_CLASS"] - )() - self._async_query_manager.init_app(app) - - def instance(self) -> AsyncQueryManager: - return self._async_query_manager diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index fc0feaf31df..1d2a8e305bc 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -29,7 +29,6 @@ from marshmallow import ValidationError from werkzeug.utils import secure_filename from superset import is_feature_enabled, security_manager -from superset.async_events.async_query_manager import AsyncQueryTokenException from superset.charts.api import ChartRestApi from superset.charts.client_processing import apply_client_processing from superset.charts.data.dashboard_filter_context import ( @@ -38,11 +37,7 @@ from superset.charts.data.dashboard_filter_context import ( get_dashboard_filter_context, ) from superset.charts.data.form_data import set_form_data -from superset.charts.data.query_context_cache_loader import QueryContextCacheLoader from superset.charts.schemas import ChartDataQueryContextSchema -from superset.commands.chart.data.create_async_job_command import ( - CreateAsyncChartDataJobCommand, -) from superset.commands.chart.data.get_data_command import ChartDataCommand from superset.commands.chart.data.streaming_export_command import ( StreamingCSVExportCommand, @@ -59,6 +54,7 @@ from superset.daos.exceptions import DatasourceNotFound from superset.exceptions import QueryObjectValidationError, SupersetSecurityException from superset.extensions import event_logger from superset.models.sql_lab import Query +from superset.tasks.async_queries import submit_chart_data_query_tasks from superset.utils import json from superset.utils.core import ( create_zip, @@ -77,7 +73,7 @@ logger = logging.getLogger(__name__) class ChartDataRestApi(ChartRestApi): - include_route_methods = {"get_data", "data", "data_from_cache"} + include_route_methods = {"get_data", "data"} @expose("//data/", methods=("GET",)) @protect() @@ -363,75 +359,6 @@ class ChartDataRestApi(ChartRestApi): expected_rows=expected_rows, ) - @expose("/data/", methods=("GET",)) - @protect() - @statsd_metrics - @event_logger.log_this_with_context( - action=lambda self, *args, **kwargs: ( - f"{self.__class__.__name__}.data_from_cache" - ), - log_to_statsd=False, - ) - def data_from_cache(self, cache_key: str) -> Response: - """ - Take a query context cache key and return payload - data response for the given query. - --- - get: - summary: Return payload data response for the given query - description: >- - Takes a query context cache key and returns payload data - response for the given query. - parameters: - - in: path - schema: - type: string - name: cache_key - responses: - 200: - description: Query result - content: - application/json: - schema: - $ref: "#/components/schemas/ChartDataResponseSchema" - 400: - $ref: '#/components/responses/400' - 401: - $ref: '#/components/responses/401' - 403: - $ref: '#/components/responses/403' - 404: - $ref: '#/components/responses/404' - 422: - $ref: '#/components/responses/422' - 500: - $ref: '#/components/responses/500' - """ - try: - cached_data = self._load_query_context_form_from_cache(cache_key) - # Set form_data in Flask Global as it is used as a fallback - # for async queries with jinja context - set_form_data(cached_data) - query_context = self._create_query_context_from_form(cached_data) - # Mark as a cache replay so _sql_filters_modified skips the - # SQL-extras check. The original request already passed the - # full security check, cache keys are opaque SHA-256 hashes - # (unguessable), and force_cached only serves pre-computed - # data — no new SQL is executed. - query_context._from_cache_replay = True - command = ChartDataCommand(query_context) - command.validate() - except ChartDataCacheLoadError: - return self.response_404() - except SupersetSecurityException: - return self.response_403() - except ValidationError as error: - return self.response_400( - message=_("Request is incorrect: %(error)s", error=error.messages) - ) - - return self._get_data_response(command, True) - def _run_async( self, form_data: dict[str, Any], @@ -454,18 +381,12 @@ class ChartDataRestApi(ChartRestApi): return self._send_chart_response(result) except ChartDataCacheLoadError: pass - # Otherwise, kick off a background job to run the chart query. - # Clients will either poll or be notified of query completion, - # at which point they will call the /data/ endpoint - # to retrieve the results. - async_command = CreateAsyncChartDataJobCommand() - try: - async_command.validate(request) - except AsyncQueryTokenException: - return self.response_401() - - async_result = async_command.run(form_data, get_user_id()) - return self.response(202, **async_result) + # Otherwise, kick off background GTF tasks (one per QueryObject) to run the + # chart query. The client polls /api/v1/task/status_changes, aggregates the + # tasks' statuses, and on success re-issues this same request — now served + # synchronously from the per-query DATA cache the tasks populated. + job = submit_chart_data_query_tasks(command.query_context, get_user_id()) + return self.response(202, **job) def _send_chart_response( # noqa: C901 self, @@ -694,10 +615,6 @@ class ChartDataRestApi(ChartRestApi): return filename, expected_rows - # pylint: disable=invalid-name - def _load_query_context_form_from_cache(self, cache_key: str) -> dict[str, Any]: - return QueryContextCacheLoader.load(cache_key) - def _map_form_data_datasource_to_dataset_id( self, form_data: dict[str, Any] ) -> dict[str, Any]: diff --git a/superset/charts/data/query_context_cache_loader.py b/superset/charts/data/query_context_cache_loader.py deleted file mode 100644 index 1bdabd33f48..00000000000 --- a/superset/charts/data/query_context_cache_loader.py +++ /dev/null @@ -1,30 +0,0 @@ -# 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. -from typing import Any - -from superset import cache -from superset.commands.chart.exceptions import ChartDataCacheLoadError - - -class QueryContextCacheLoader: # pylint: disable=too-few-public-methods - @staticmethod - def load(cache_key: str) -> dict[str, Any]: - cache_value = cache.get(cache_key) - if not cache_value: - raise ChartDataCacheLoadError("Cached data not found") - - return cache_value["data"] diff --git a/superset/commands/chart/data/create_async_job_command.py b/superset/commands/chart/data/create_async_job_command.py deleted file mode 100644 index a212308d104..00000000000 --- a/superset/commands/chart/data/create_async_job_command.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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 logging -from typing import Any, Optional - -from flask import Request - -from superset.extensions import async_query_manager - -logger = logging.getLogger(__name__) - - -class CreateAsyncChartDataJobCommand: - _async_channel_id: str - - def validate(self, request: Request) -> None: - self._async_channel_id = async_query_manager.parse_channel_id_from_request( - request - ) - - def run(self, form_data: dict[str, Any], user_id: Optional[int]) -> dict[str, Any]: - if not getattr(self, "_async_channel_id", None): - raise RuntimeError( - "CreateAsyncChartDataJobCommand.run() called before validate(); " - "the async channel id was not initialized." - ) - return async_query_manager.submit_chart_data_job( - self._async_channel_id, form_data, user_id - ) diff --git a/superset/commands/chart/data/get_data_command.py b/superset/commands/chart/data/get_data_command.py index 946811269bf..7c4942b0262 100644 --- a/superset/commands/chart/data/get_data_command.py +++ b/superset/commands/chart/data/get_data_command.py @@ -38,6 +38,10 @@ class ChartDataCommand(BaseCommand): def __init__(self, query_context: QueryContext): self._query_context = query_context + @property + def query_context(self) -> QueryContext: + return self._query_context + def run(self, **kwargs: Any) -> dict[str, Any]: # caching is handled in query_context.get_df_payload # (also evals `force` property) diff --git a/superset/commands/tasks/cancel.py b/superset/commands/tasks/cancel.py index cbc07dc889e..2416f071dc7 100644 --- a/superset/commands/tasks/cancel.py +++ b/superset/commands/tasks/cancel.py @@ -188,8 +188,16 @@ class CancelTaskCommand(BaseCommand): ) if task.is_shared: - # Shared tasks: must be a subscriber - if not user_id or not task.has_subscriber(user_id): + # Shared tasks: must be a subscriber. Embedded guests have no user_id + # and subscribe by a token-derived guest_key instead (see + # superset.tasks.guest), so honor either identity. + from superset.tasks.guest import get_current_guest_subscriber_key + + guest_key = None if user_id else get_current_guest_subscriber_key() + subscribed = (user_id and task.has_subscriber(user_id)) or ( + guest_key and task.has_guest_subscriber(guest_key) + ) + if not subscribed: raise TaskPermissionDeniedError( "You must be subscribed to cancel this shared task" ) @@ -267,23 +275,32 @@ class CancelTaskCommand(BaseCommand): def _do_unsubscribe(self, task: "Task", user_id: int | None) -> "Task": """ - Execute unsubscribe operation. + Execute unsubscribe operation (user or embedded guest). :param task: The task to unsubscribe from - :param user_id: ID of user to unsubscribe + :param user_id: ID of user to unsubscribe, or None for an embedded guest :returns: The updated task model """ from superset.daos.tasks import TaskDAO + from superset.tasks.guest import get_current_guest_subscriber_key self._action_taken = "unsubscribed" - if not user_id or not task.has_subscriber(user_id): - # User not subscribed - they shouldn't be able to cancel + # Embedded guests subscribe by a token-derived key, not a user_id. + guest_key = None if user_id else get_current_guest_subscriber_key() + + if user_id and task.has_subscriber(user_id): + result = TaskDAO.remove_subscriber(task.id, user_id) + subscriber = f"user {user_id}" + elif guest_key and task.has_guest_subscriber(guest_key): + result = TaskDAO.remove_guest_subscriber(task.id, guest_key) + subscriber = "guest" + else: + # Not subscribed - they shouldn't be able to cancel raise TaskPermissionDeniedError( "You are not subscribed to this shared task" ) - result = TaskDAO.remove_subscriber(task.id, user_id) if result is None: raise TaskPermissionDeniedError( "You are not subscribed to this shared task" @@ -294,8 +311,8 @@ class CancelTaskCommand(BaseCommand): stats_logger.incr("gtf.task.unsubscribe") logger.info( - "User %s unsubscribed from shared task: %s", - user_id, + "%s unsubscribed from shared task: %s", + subscriber, task.uuid, ) diff --git a/superset/commands/tasks/submit.py b/superset/commands/tasks/submit.py index f3e627ce325..c2b89881d1c 100644 --- a/superset/commands/tasks/submit.py +++ b/superset/commands/tasks/submit.py @@ -34,6 +34,7 @@ from superset.commands.tasks.exceptions import ( ) from superset.daos.exceptions import DAOCreateFailedError from superset.stats_logger import BaseStatsLogger +from superset.tasks.guest import get_current_guest_subscriber_key from superset.tasks.locks import task_lock from superset.tasks.utils import get_active_dedup_key from superset.utils.core import get_user_id @@ -92,6 +93,9 @@ class SubmitTaskCommand(BaseCommand): task_key = self._properties.get("task_key") or str(uuid.uuid4()) scope = self._properties.get("scope", TaskScope.PRIVATE.value) user_id = get_user_id() + # Embedded guests have no ab_user id; they subscribe by a token-derived + # key so TaskFilter can grant them visibility of their own tasks. + guest_key = None if user_id else get_current_guest_subscriber_key() # Build dedup_key for lock dedup_key = get_active_dedup_key( @@ -119,6 +123,11 @@ class SubmitTaskCommand(BaseCommand): user_id, task_key, ) + elif guest_key and not existing.has_guest_subscriber(guest_key): + # Embedded guest joining a SHARED task an equivalent guest + # created; subscribe so this guest can also poll it. + TaskDAO.add_guest_subscriber(existing.id, guest_key) + stats_logger.incr("gtf.task.subscribe") else: # Same user submitted the same task - deduplication hit stats_logger.incr("gtf.task.dedupe") @@ -137,6 +146,7 @@ class SubmitTaskCommand(BaseCommand): scope=scope, task_name=self._properties.get("task_name"), user_id=user_id, + guest_key=guest_key, payload=self._properties.get("payload", {}), properties=self._properties.get("properties", {}), ) diff --git a/superset/common/query_context.py b/superset/common/query_context.py index 8a5819e6aa6..d73669dd8fc 100644 --- a/superset/common/query_context.py +++ b/superset/common/query_context.py @@ -133,6 +133,17 @@ class QueryContext: def query_cache_key(self, query_obj: QueryObject, **kwargs: Any) -> str | None: return self._processor.query_cache_key(query_obj, **kwargs) + def prepare_contribution_totals(self) -> tuple[list[int], int | None]: + """Identify contribution queries and normalize the totals query. + + Returns the indices of queries whose contribution post-processing needs a + shared totals row, and the index of the totals query itself (or ``None``). + As a side effect the totals query's ``row_limit`` is cleared so its cache + key matches the entry its dependents read — see the synchronous equivalent + in ``QueryContextProcessor.get_payload_result``. + """ + return self._processor._prepare_contribution_totals() + def get_df_payload( self, query_obj: QueryObject, diff --git a/superset/config.py b/superset/config.py index 9bc99696c7e..b8c26ef325c 100644 --- a/superset/config.py +++ b/superset/config.py @@ -53,7 +53,6 @@ from superset.advanced_data_type.plugins.internet_address import internet_addres from superset.advanced_data_type.plugins.internet_port import internet_port from superset.advanced_data_type.types import AdvancedDataType from superset.constants import ( - CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET, CHANGE_ME_GUEST_TOKEN_JWT_SECRET, CHANGE_ME_SECRET_KEY, ) @@ -2914,63 +2913,13 @@ SQLA_TABLE_MUTATOR = lambda table: table # noqa: E731 # Global async query config options. -# Requires GLOBAL_ASYNC_QUERIES feature flag to be enabled. -GLOBAL_ASYNC_QUERY_MANAGER_CLASS = ( - "superset.async_events.async_query_manager.AsyncQueryManager" -) -GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX = "async-events-" -GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT = 1000 -GLOBAL_ASYNC_QUERIES_REDIS_STREAM_LIMIT_FIREHOSE = 1000000 -GLOBAL_ASYNC_QUERIES_REGISTER_REQUEST_HANDLERS = True -GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME = "async-token" -GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SECURE = False -GLOBAL_ASYNC_QUERIES_JWT_COOKIE_SAMESITE: None | (Literal["None", "Lax", "Strict"]) = ( - None -) -GLOBAL_ASYNC_QUERIES_JWT_COOKIE_DOMAIN = None -GLOBAL_ASYNC_QUERIES_JWT_SECRET = CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET -# Lifetime of the async-query JWT, in seconds. After this period the token -# expires and a fresh one is issued on the next request. -GLOBAL_ASYNC_QUERIES_JWT_EXPIRATION_SECONDS = int(timedelta(hours=1).total_seconds()) -GLOBAL_ASYNC_QUERIES_TRANSPORT: Literal["polling", "ws"] = "polling" +# Requires the GLOBAL_ASYNC_QUERIES feature flag to be enabled. Async chart-data +# queries run on the Global Task Framework (one task per QueryObject) over +# DISTRIBUTED_COORDINATION_CONFIG; the client polls /api/v1/task/status_changes at +# this interval (milliseconds) and re-issues its request once the tasks succeed. GLOBAL_ASYNC_QUERIES_POLLING_DELAY = int( timedelta(milliseconds=500).total_seconds() * 1000 ) -GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL = "ws://127.0.0.1:8080/" - -# Global async queries cache backend configuration options: -# - Set 'CACHE_TYPE' to 'RedisCache' for RedisCacheBackend. -# - Set 'CACHE_TYPE' to 'RedisSentinelCache' for RedisSentinelCacheBackend. -# -# DEPRECATED: this dedicated backend is retained only so Global Async Queries can -# keep running on their own coordination connection when DISTRIBUTED_COORDINATION_CONFIG -# is not configured. When configured it is used by GAQ *only* (its event streams and -# cancel registry); it never powers distributed locks or the Global Task Framework, -# which use DISTRIBUTED_COORDINATION_CONFIG exclusively. GAQ uses -# DISTRIBUTED_COORDINATION_CONFIG whenever it is set and only falls back to this -# dedicated backend when it is not, so a consolidated deployment need not maintain a -# second config. This dual-backend arrangement is deprecated and removed in Superset -# 8.0, when GAQ moves onto DISTRIBUTED_COORDINATION_CONFIG like the rest of Superset's -# coordination. All parameters here are supported identically under -# DISTRIBUTED_COORDINATION_CONFIG (both use the same -# RedisCache/RedisSentinelCache backend). -GLOBAL_ASYNC_QUERIES_CACHE_BACKEND = { - "CACHE_TYPE": "RedisCache", - "CACHE_REDIS_HOST": "localhost", - "CACHE_REDIS_PORT": 6379, - "CACHE_REDIS_USER": "", - "CACHE_REDIS_PASSWORD": "", - "CACHE_REDIS_DB": 0, - "CACHE_DEFAULT_TIMEOUT": 300, - "CACHE_REDIS_SENTINELS": [("localhost", 26379)], - "CACHE_REDIS_SENTINEL_MASTER": "mymaster", - "CACHE_REDIS_SENTINEL_PASSWORD": None, - "CACHE_REDIS_SSL": False, # True or False - "CACHE_REDIS_SSL_CERTFILE": None, - "CACHE_REDIS_SSL_KEYFILE": None, - "CACHE_REDIS_SSL_CERT_REQS": "required", - "CACHE_REDIS_SSL_CA_CERTS": None, -} # Embedded config options GUEST_ROLE_NAME = "Public" @@ -3236,31 +3185,19 @@ TASK_PROGRESS_UPDATE_THROTTLE_INTERVAL = 2 # seconds # These features require Redis primitives unavailable in generic cache backends: # - Pub/Sub: Real-time message broadcasting between workers # - SET NX EX: Atomic lock acquisition with automatic expiration -# - Streams: Persistent ordered event logs (e.g. the Global Async Queries firehose) +# - Streams: Persistent ordered event logs (task completion signalling) # # When configured, enables: # - Real-time abort/completion notifications for GTF tasks (vs database polling) # - Redis-based distributed locking (vs KeyValueDAO-backed DistributedLock) -# - Global Async Queries event streams (the async-events / firehose transport) +# - Async chart-data queries (Global Task Framework task streams) # # This backend powers the higher-level coordination service # (``superset.coordination.base.CoordinationService``) exposing standardized interfaces # for distributed locks, pub/sub, and streams under a single connection. It is the -# single source of truth for the coordinator's consumers (distributed locks, the -# Global Task Framework, and future stream/pub-sub users). Global Async Queries use -# this connection whenever it is set, falling back to their dedicated -# ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` only when it is unset; that dual-backend -# arrangement is deprecated and, in Superset 8.0, GAQ moves onto this connection and -# the dedicated backend is removed. -# -# All parameters previously supported by GLOBAL_ASYNC_QUERIES_CACHE_BACKEND are -# supported here (both go through the same RedisCacheBackend/RedisSentinelCacheBackend -# `from_config`): CACHE_REDIS_HOST, CACHE_REDIS_PORT, CACHE_REDIS_USER, -# CACHE_REDIS_PASSWORD, CACHE_REDIS_DB, CACHE_KEY_PREFIX, CACHE_DEFAULT_TIMEOUT, -# CACHE_REDIS_SSL, CACHE_REDIS_SSL_CERTFILE, CACHE_REDIS_SSL_KEYFILE, -# CACHE_REDIS_SSL_CERT_REQS, CACHE_REDIS_SSL_CA_CERTS, CACHE_REDIS_SOCKET_TIMEOUT, -# CACHE_REDIS_SOCKET_CONNECT_TIMEOUT, and for Sentinel CACHE_REDIS_SENTINELS, -# CACHE_REDIS_SENTINEL_MASTER, CACHE_REDIS_SENTINEL_PASSWORD. +# single source of truth for the coordinator's consumers: distributed locks, the +# Global Task Framework (including async chart-data queries), and future +# stream/pub-sub users. # # Example with standard Redis: # DISTRIBUTED_COORDINATION_CONFIG: CacheConfig = { diff --git a/superset/constants.py b/superset/constants.py index 3ceede47bdb..77e8d17cf26 100644 --- a/superset/constants.py +++ b/superset/constants.py @@ -29,7 +29,6 @@ EMPTY_STRING = "" CHANGE_ME_SECRET_KEY = "CHANGE_ME_TO_A_COMPLEX_RANDOM_SECRET" # noqa: S105 CHANGE_ME_GUEST_TOKEN_JWT_SECRET = "test-guest-secret-change-me" # noqa: S105 -CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET = "test-secret-change-me" # noqa: S105 SKIP_VISIBILITY_FILTER_CLASSES = "_skip_visibility_filter_classes" diff --git a/superset/coordination/__init__.py b/superset/coordination/__init__.py index 9fb3958661a..a259c4088e7 100644 --- a/superset/coordination/__init__.py +++ b/superset/coordination/__init__.py @@ -25,12 +25,12 @@ backend when one is configured and falls back to a database-backed lock otherwis Historically these were wired up independently: the Global Task Framework used ``DISTRIBUTED_COORDINATION_CONFIG`` (pub/sub and locking) while Global Async Queries -used a separate ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` for its streams, and each caller -hand-rolled its own pub/sub-vs-poll wait loops. Consolidating them here keeps the -architecture modular, gives other components (e.g. the extensions framework) a single -reusable coordination surface, and reduces the number of moving parts. The legacy -``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND`` is still honored as a fallback (with a -deprecation warning) so existing deployments keep working during the transition. +used a separate cache backend for their event streams, and each caller hand-rolled +its own pub/sub-vs-poll wait loops. Consolidating them here keeps the architecture +modular, gives other components (e.g. the extensions framework) a single reusable +coordination surface, and reduces the number of moving parts. All coordination — +including async chart-data queries, which now run on the Global Task Framework — +uses ``DISTRIBUTED_COORDINATION_CONFIG`` exclusively. Import concrete classes directly from their modules: :class:`~superset.coordination.base.CoordinationService`, diff --git a/superset/coordination/base.py b/superset/coordination/base.py index 5c081e91546..7a1c2361e70 100644 --- a/superset/coordination/base.py +++ b/superset/coordination/base.py @@ -93,13 +93,9 @@ class CoordinationService: Returns the shared coordination connection (via the cache manager), or ``None`` when ``DISTRIBUTED_COORDINATION_CONFIG`` is not configured. This is - the single source of truth for the coordinator's consumers (distributed - locks, the Global Task Framework, and future stream/pub-sub users); it does - *not* consult the deprecated ``GLOBAL_ASYNC_QUERIES_CACHE_BACKEND``. Global - Async Queries resolve their own backend (this coordinator when configured, - else the deprecated dedicated backend — see - :class:`~superset.async_events.async_query_manager.AsyncQueryManager`) and pass - it explicitly to the primitives below via ``backend``. + the single source of truth for the coordinator's consumers: distributed + locks, the Global Task Framework (including async chart-data queries), and + future stream/pub-sub users. """ from superset.extensions import cache_manager diff --git a/superset/daos/tasks.py b/superset/daos/tasks.py index 41001ac94ea..59cc587b387 100644 --- a/superset/daos/tasks.py +++ b/superset/daos/tasks.py @@ -68,6 +68,58 @@ class TaskDAO(BaseDAO[Task]): result = query.with_entities(Task.status).one_or_none() return result[0] if result else None + @classmethod + def get_statuses_changed_since( + cls, cursor: datetime | None, task_type: str | None = None + ) -> tuple[dict[str, dict[str, Any]], datetime]: + """Return ``{uuid: {status, progress}}`` for tasks changed since ``cursor``. + + The minimal-IO polling primitive behind both the async chart-data + completion poll and the realtime task list. The base filter + (``TaskFilter``) scopes results to tasks the caller can see (subscribed + tasks for regular users, all tasks for admins), so callers never pass an + explicit id list. ``task_type`` optionally narrows to a single kind (e.g. + chart-data query tasks) so a client tracks only the work it cares about. + + Without a ``cursor`` this establishes a **baseline**: it returns no + statuses and a fresh watermark (the current server clock), so a client + gets a definitive starting point without dumping every task in the + metastore. Subsequent calls pass that watermark back and receive only + tasks whose ``changed_on >= cursor`` (``>=``, not ``>``, so no transition + straddling the boundary is missed — re-delivery of an already-seen status + is idempotent for the client, a miss would hang it). + + Returns the ``{uuid: {status, progress}}`` map (``progress`` is the + 0.0–1.0 percent from the task's properties, or ``None`` when unknown) plus + the next cursor to poll with (the max ``changed_on`` in this batch), so the + client always advances using a server-observed watermark, never its clock. + """ + # Baseline: no cursor → start "from now", surfacing only later changes. + if cursor is None: + return {}, datetime.now() + + query = cls._apply_base_filter(db.session.query(Task)).filter( + # Task.changed_on's type is shadowed by CoreTask's bare annotation + # (datetime | None), so reference the real column for the comparison. + Task.__table__.c.changed_on >= cursor + ) + if task_type is not None: + query = query.filter(Task.task_type == task_type) + rows = query.with_entities( + Task.uuid, Task.status, Task.changed_on, Task.properties + ).all() + + statuses: dict[str, dict[str, Any]] = {} + changed_times: list[datetime] = [] + for uuid, status, changed_on, properties in rows: + progress = json.loads(properties or "{}").get("progress_percent") + statuses[str(uuid)] = {"status": status, "progress": progress} + if changed_on is not None: + changed_times.append(changed_on) + # Advance to the newest change seen, or hold the cursor if nothing changed. + next_cursor = max(changed_times, default=cursor) + return statuses, next_cursor + @classmethod def find_by_task_key( cls, @@ -109,6 +161,7 @@ class TaskDAO(BaseDAO[Task]): task_key: str, scope: TaskScope | str = TaskScope.PRIVATE, user_id: int | None = None, + guest_key: str | None = None, payload: dict[str, Any] | None = None, properties: TaskProperties | None = None, **kwargs: Any, @@ -181,6 +234,10 @@ class TaskDAO(BaseDAO[Task]): task_key, scope_value, ) + elif guest_key: + # Embedded guest creator: subscribe by token-derived key so the guest + # can see the task it just created (see superset.tasks.guest). + cls.add_guest_subscriber(task.id, guest_key) logger.info( "Created new async task: %s (type: %s, scope: %s)", @@ -289,6 +346,39 @@ class TaskDAO(BaseDAO[Task]): logger.info("Added subscriber %s to task %s", user_id, task_id) return True + @classmethod + def add_guest_subscriber(cls, task_id: int, guest_key: str) -> bool: + """ + Subscribe an embedded guest (by token-derived key) to a task. + + The guest counterpart of ``add_subscriber``: guests have no ``ab_user`` + row, so they subscribe by ``guest_key`` (see ``superset.tasks.guest``), + which grants them visibility of the task through ``TaskFilter``. + + :param task_id: ID of the task + :param guest_key: Stable guest identity to subscribe + :returns: True if subscriber was added, False if already exists + """ + # Check first to avoid IntegrityError (unrecoverable in nested txns). + existing = ( + db.session.query(TaskSubscriber) + .filter_by(task_id=task_id, guest_key=guest_key) + .first() + ) + if existing: + return False + + db.session.add( + TaskSubscriber( + task_id=task_id, + guest_key=guest_key, + subscribed_at=datetime.now(timezone.utc), + ) + ) + db.session.flush() + logger.info("Added guest subscriber to task %s", task_id) + return True + @classmethod def remove_subscriber(cls, task_id: int, user_id: int) -> Task | None: """ @@ -303,12 +393,33 @@ class TaskDAO(BaseDAO[Task]): :returns: Updated Task if subscriber was removed, None if not subscribed :raises DAODeleteFailedError: If subscription removal fails """ + return cls._remove_subscription( + task_id, TaskSubscriber.user_id == user_id, f"user {user_id}" + ) + + @classmethod + def remove_guest_subscriber(cls, task_id: int, guest_key: str) -> Task | None: + """ + Remove an embedded guest's subscription (by ``guest_key``) from a task. + + The guest counterpart of ``remove_subscriber`` (see ``superset.tasks.guest``). + + :param task_id: ID of the task + :param guest_key: Guest identity to unsubscribe + :returns: Updated Task if subscriber was removed, None if not subscribed + :raises DAODeleteFailedError: If subscription removal fails + """ + return cls._remove_subscription( + task_id, TaskSubscriber.guest_key == guest_key, "guest" + ) + + @classmethod + def _remove_subscription( + cls, task_id: int, subscriber_clause: Any, label: str + ) -> Task | None: subscription = ( db.session.query(TaskSubscriber) - .filter( - TaskSubscriber.task_id == task_id, - TaskSubscriber.user_id == user_id, - ) + .filter(TaskSubscriber.task_id == task_id, subscriber_clause) .one_or_none() ) @@ -318,7 +429,7 @@ class TaskDAO(BaseDAO[Task]): try: db.session.delete(subscription) db.session.flush() - logger.info("Removed subscriber %s from task %s", user_id, task_id) + logger.info("Removed subscriber %s from task %s", label, task_id) # Return the updated task task = cls.find_by_id(task_id, skip_base_filter=True) @@ -330,7 +441,7 @@ class TaskDAO(BaseDAO[Task]): raise except Exception as ex: raise DAODeleteFailedError( - f"Failed to remove subscription for task {task_id}, user {user_id}" + f"Failed to remove subscription for task {task_id} ({label})" ) from ex # Dependency (DAG) management methods diff --git a/superset/extensions/__init__.py b/superset/extensions/__init__.py index b5035b73c24..b1cef9c3ad8 100644 --- a/superset/extensions/__init__.py +++ b/superset/extensions/__init__.py @@ -40,8 +40,6 @@ from flask_talisman import Talisman from flask_wtf.csrf import CSRFProtect from werkzeug.local import LocalProxy -from superset.async_events.async_query_manager import AsyncQueryManager -from superset.async_events.async_query_manager_factory import AsyncQueryManagerFactory from superset.extensions.ssh import SSHManagerFactory from superset.extensions.stats_logger import BaseStatsLoggerManager from superset.security.manager import SupersetSecurityManager @@ -147,10 +145,6 @@ class ProfilingExtension: # pylint: disable=too-few-public-methods APP_DIR = os.path.join(os.path.dirname(__file__), os.path.pardir) appbuilder = AppBuilder(update_perms=False) -async_query_manager_factory = AsyncQueryManagerFactory() -async_query_manager: AsyncQueryManager = LocalProxy( - async_query_manager_factory.instance -) cache_manager = CacheManager() celery_app = celery.Celery() csrf = CSRFProtect() diff --git a/superset/initialization/__init__.py b/superset/initialization/__init__.py index c37734bb58c..59d3630a416 100644 --- a/superset/initialization/__init__.py +++ b/superset/initialization/__init__.py @@ -50,7 +50,6 @@ from werkzeug.middleware.proxy_fix import ProxyFix from superset.commands.database.exceptions import DatabaseInvalidError from superset.constants import ( - CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET, CHANGE_ME_GUEST_TOKEN_JWT_SECRET, CHANGE_ME_SECRET_KEY, ) @@ -59,7 +58,6 @@ from superset.extensions import ( _event_logger, APP_DIR, appbuilder, - async_query_manager_factory, cache_manager, celery_app, csrf, @@ -181,7 +179,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods from superset.advanced_data_type.api import AdvancedDataTypeRestApi from superset.annotation_layers.annotations.api import AnnotationRestApi from superset.annotation_layers.api import AnnotationLayerRestApi - from superset.async_events.api import AsyncEventsRestApi from superset.available_domains.api import AvailableDomainsRestApi from superset.cachekeys.api import CacheRestApi from superset.charts.api import ChartRestApi @@ -271,7 +268,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods # appbuilder.add_api(AnnotationRestApi) appbuilder.add_api(AnnotationLayerRestApi) - appbuilder.add_api(AsyncEventsRestApi) appbuilder.add_api(AdvancedDataTypeRestApi) appbuilder.add_api(AvailableDomainsRestApi) appbuilder.add_api(CacheRestApi) @@ -1047,7 +1043,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods self.configure_url_map_converters() self.configure_data_sources() self.configure_auth_provider() - self.configure_async_queries() self.configure_ssh_manager() self.configure_stats_manager() self.configure_task_manager() @@ -1138,32 +1133,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods ) sys.exit(1) - def check_async_query_secret(self) -> None: - """Refuse to start with the default async JWT secret when GAQ is enabled.""" - if not feature_flag_manager.is_feature_enabled("GLOBAL_ASYNC_QUERIES"): - return - if ( - self.config.get("GLOBAL_ASYNC_QUERIES_JWT_SECRET") - != CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET - ): - return - self._log_config_warning( - "GLOBAL_ASYNC_QUERIES is enabled but GLOBAL_ASYNC_QUERIES_JWT_SECRET " - "has not been changed from its default value.\n" - "The default value is publicly known and must be replaced before " - "running in production.\n" - "Set a strong random value (at least 32 bytes) in superset_config.py:\n" - " GLOBAL_ASYNC_QUERIES_JWT_SECRET = " - "''" - ) - if self.superset_app.debug or self.superset_app.config["TESTING"] or is_test(): - return - logger.error( - "Refusing to start: insecure GLOBAL_ASYNC_QUERIES_JWT_SECRET " - "with GLOBAL_ASYNC_QUERIES enabled" - ) - sys.exit(1) - def check_encryption_engine(self) -> None: """Warn when app-encrypted fields use the legacy AES-CBC engine. @@ -1347,7 +1316,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods # conditionally self.configure_feature_flags() self.check_guest_token_secret() - self.check_async_query_secret() self.check_encryption_engine() self.configure_db_encrypt() self.setup_db() @@ -1629,20 +1597,6 @@ class SupersetAppInitializer: # pylint: disable=too-many-public-methods for ex in csrf_exempt_list: csrf.exempt(ex) - def configure_async_queries(self) -> None: - if feature_flag_manager.is_feature_enabled("GLOBAL_ASYNC_QUERIES"): - # In production, check_async_query_secret() already aborts startup when - # the default secret is present, so this branch is never reached with it. - # In debug/testing the check only warns, so skip async-query init here to - # avoid AsyncQueryManager.init_app() hard-failing on the too-short default - # secret and crashing startup despite the warn-only intent. - if ( - self.config.get("GLOBAL_ASYNC_QUERIES_JWT_SECRET") - == CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET - ): - return - async_query_manager_factory.init_app(self.superset_app) - def configure_task_manager(self) -> None: """Initialize the TaskManager for GTF realtime notifications.""" if feature_flag_manager.is_feature_enabled("GLOBAL_TASK_FRAMEWORK"): diff --git a/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py b/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py index 6e9b08b0fca..9c70c3aac2c 100644 --- a/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py +++ b/superset/migrations/versions/2026-08-21_12-00_7e2c9a4f1b83_create_task_dependencies_table.py @@ -14,7 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -"""Create task_dependencies table for Global Task Framework (GTF) task DAG +"""Create task_dependencies table and add task_subscribers.guest_key (GTF) Revision ID: 7e2c9a4f1b83 Revises: 1072de5ed955 @@ -22,6 +22,8 @@ Create Date: 2026-08-21 12:00:00.000000 """ +import sqlalchemy as sa +from alembic import op from sqlalchemy import ( Column, DateTime, @@ -30,9 +32,11 @@ from sqlalchemy import ( ) from superset.migrations.shared.utils import ( + add_columns, create_fks_for_table, create_index, create_table, + drop_columns, drop_fks_for_table, drop_index, drop_table, @@ -44,17 +48,23 @@ down_revision = "1072de5ed955" TASKS_TABLE = "tasks" TASK_DEPENDENCIES_TABLE = "task_dependencies" +TASK_SUBSCRIBERS_TABLE = "task_subscribers" def upgrade(): """ - Create the task_dependencies junction table for the task dependency graph. + Create the task_dependencies junction table and add task_subscribers.guest_key. - Each row is a directed edge: ``task_id`` (the dependent) depends on - ``depends_on_task_id`` (the prerequisite). Both foreign keys reference - ``tasks.id`` with ``ON DELETE CASCADE`` so edges are removed when either - endpoint task is pruned (task pruning uses a bulk core DELETE that bypasses - the ORM cascade, so the database-level cascade is required for cleanup). + ``task_dependencies``: each row is a directed edge — ``task_id`` (the + dependent) depends on ``depends_on_task_id`` (the prerequisite). Both foreign + keys reference ``tasks.id`` with ``ON DELETE CASCADE`` so edges are removed + when either endpoint task is pruned (task pruning uses a bulk core DELETE that + bypasses the ORM cascade, so the database-level cascade is required for + cleanup). + + ``task_subscribers.guest_key``: lets embedded guests (which have no + ``ab_user`` row) subscribe to tasks by a stable, token-derived key so the task + filter can grant them visibility of their own async work. """ create_table( TASK_DEPENDENCIES_TABLE, @@ -119,9 +129,37 @@ def upgrade(): ondelete="SET NULL", ) + # Let embedded guests subscribe to tasks by a token-derived ``guest_key``. + # Guests have no ``ab_user`` row, so a subscription is identified by exactly + # one of ``user_id`` (authenticated) or ``guest_key`` (guest): add the + # nullable ``guest_key`` column, relax ``user_id`` to nullable, and add a + # unique ``(task_id, guest_key)`` index mirroring the existing + # ``(task_id, user_id)`` uniqueness so a guest subscribes at most once. + # (NULLs are distinct in unique constraints, so user rows and guest rows do + # not collide.) + add_columns( + TASK_SUBSCRIBERS_TABLE, + Column("guest_key", sa.String(length=64), nullable=True), + ) + with op.batch_alter_table(TASK_SUBSCRIBERS_TABLE) as batch_op: + batch_op.alter_column("user_id", existing_type=sa.Integer(), nullable=True) + batch_op.create_index("ix_task_subscribers_guest_key", ["guest_key"]) + batch_op.create_unique_constraint( + "uq_task_subscribers_task_guest", ["task_id", "guest_key"] + ) + def downgrade(): - """Drop the task_dependencies table and its indexes and foreign keys.""" + """Drop task_dependencies and revert the task_subscribers.guest_key change.""" + # Guest subscriptions cannot be represented without the column; drop those + # rows first so restoring user_id NOT NULL does not fail on NULL user_id. + op.execute(sa.text("DELETE FROM task_subscribers WHERE user_id IS NULL")) + with op.batch_alter_table(TASK_SUBSCRIBERS_TABLE) as batch_op: + batch_op.drop_constraint("uq_task_subscribers_task_guest", type_="unique") + batch_op.drop_index("ix_task_subscribers_guest_key") + batch_op.alter_column("user_id", existing_type=sa.Integer(), nullable=False) + drop_columns(TASK_SUBSCRIBERS_TABLE, "guest_key") + drop_fks_for_table( TASK_DEPENDENCIES_TABLE, [ diff --git a/superset/models/task_subscribers.py b/superset/models/task_subscribers.py index d069c7f3b00..f56f4621447 100644 --- a/superset/models/task_subscribers.py +++ b/superset/models/task_subscribers.py @@ -19,7 +19,14 @@ from datetime import datetime, timezone from flask_appbuilder import Model -from sqlalchemy import Column, DateTime, ForeignKey, Integer, UniqueConstraint +from sqlalchemy import ( + Column, + DateTime, + ForeignKey, + Integer, + String, + UniqueConstraint, +) from sqlalchemy.orm import relationship from superset_core.tasks.models import TaskSubscriber as CoreTaskSubscriber @@ -37,6 +44,13 @@ class TaskSubscriber(CoreTaskSubscriber, AuditMixinNullable, Model): Subscribers can unsubscribe from shared tasks. When the last subscriber unsubscribes, the task is automatically aborted. + + A subscriber is identified by exactly one of ``user_id`` (an authenticated + ``ab_user``) or ``guest_key`` (a stable, unguessable identity derived from an + embedded guest token — see ``superset.tasks.guest``). Guests have no + ``ab_user`` row, so they subscribe by ``guest_key`` and gain visibility of + the tasks they created/joined (which SHARED-scope dedup may collapse across + equivalent guests) through the same subscription mechanism as users. """ __tablename__ = "task_subscribers" @@ -45,9 +59,11 @@ class TaskSubscriber(CoreTaskSubscriber, AuditMixinNullable, Model): task_id = Column( Integer, ForeignKey("tasks.id", ondelete="CASCADE"), nullable=False ) + # Exactly one of user_id / guest_key is set. user_id = Column( - Integer, ForeignKey("ab_user.id", ondelete="CASCADE"), nullable=False + Integer, ForeignKey("ab_user.id", ondelete="CASCADE"), nullable=True ) + guest_key = Column(String(64), nullable=True, index=True) subscribed_at = Column(DateTime, nullable=False, default=datetime.now(timezone.utc)) # Relationships @@ -56,7 +72,9 @@ class TaskSubscriber(CoreTaskSubscriber, AuditMixinNullable, Model): __table_args__ = ( UniqueConstraint("task_id", "user_id", name="uq_task_subscribers_task_user"), + UniqueConstraint("task_id", "guest_key", name="uq_task_subscribers_task_guest"), ) def __repr__(self) -> str: - return f"" + subscriber = self.user_id if self.user_id is not None else self.guest_key + return f"" diff --git a/superset/models/tasks.py b/superset/models/tasks.py index 89b5da3c58a..5707373d1e1 100644 --- a/superset/models/tasks.py +++ b/superset/models/tasks.py @@ -338,13 +338,22 @@ class Task(CoreTask, AuditMixinNullable, Model): """ return any(sub.user_id == user_id for sub in self.subscribers) + def has_guest_subscriber(self, guest_key: str) -> bool: + """ + Check if an embedded guest (by token-derived key) is subscribed. + + :param guest_key: Guest identity to check (see superset.tasks.guest) + :returns: True if the guest is subscribed + """ + return any(sub.guest_key == guest_key for sub in self.subscribers) + def get_subscriber_ids(self) -> list[int]: """ Get list of all subscriber user IDs. :returns: List of user IDs subscribed to this task """ - return [sub.user_id for sub in self.subscribers] + return [sub.user_id for sub in self.subscribers if sub.user_id is not None] def to_dict(self) -> dict[str, Any]: """ diff --git a/superset/tasks/api.py b/superset/tasks/api.py index c427c74abb2..cd092379efe 100644 --- a/superset/tasks/api.py +++ b/superset/tasks/api.py @@ -17,9 +17,10 @@ """Task REST API""" import logging +from datetime import datetime from uuid import UUID -from flask import Response +from flask import request, Response from flask_appbuilder.api import expose, protect, safe from flask_appbuilder.models.sqla.interface import SQLAInterface @@ -66,6 +67,7 @@ class TaskRestApi(BaseSupersetModelRestApi): **MODEL_API_RW_METHOD_PERMISSION_MAP, "cancel": "write", "status": "read", + "status_changes": "read", } # Only allow read operations - no create/update/delete through REST API @@ -76,6 +78,7 @@ class TaskRestApi(BaseSupersetModelRestApi): RouteMethod.INFO, "cancel", "status", + "status_changes", "related_subscribers", "related", } @@ -257,6 +260,91 @@ class TaskRestApi(BaseSupersetModelRestApi): except (ValueError, TypeError): return self.response_404() + @expose("/status_changes", methods=("GET",)) + @protect() + @safe + @statsd_metrics + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" + ".status_changes", + log_to_statsd=False, + ) + def status_changes(self) -> Response: + """Poll status of accessible tasks changed since a cursor. + --- + get: + summary: Poll changed task statuses + description: > + Minimal-IO polling primitive for clients awaiting async work (e.g. + chart-data query tasks) or rendering a live task list. Returns a + ``{uuid: {status, progress}}`` map for tasks the caller can access + (subscribed tasks for regular users, all tasks for admins) that changed + since the ``cursor``, plus the next cursor to poll with. Omit ``cursor`` + on the first call to establish a baseline. Pass ``task_type`` to track + only one kind of task. + parameters: + - in: query + name: cursor + schema: + type: string + required: false + description: > + Opaque watermark returned by a previous call. Omit to establish a + baseline (returns the current watermark and no statuses). + - in: query + name: task_type + schema: + type: string + required: false + description: > + Restrict results to a single task type (e.g. the chart-data query + task type), so a client tracks only the work it cares about. + responses: + 200: + description: Changed task statuses since the cursor + content: + application/json: + schema: + type: object + properties: + statuses: + type: object + additionalProperties: + type: object + properties: + status: + type: string + progress: + type: number + nullable: true + description: Map of task UUID to status and progress + cursor: + type: string + nullable: true + description: Cursor to pass to the next poll + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + """ + from superset.daos.tasks import TaskDAO + + cursor: datetime | None = None + if cursor_arg := request.args.get("cursor"): + try: + cursor = datetime.fromisoformat(cursor_arg) + except ValueError: + return self.response_400(message="Invalid cursor") + + statuses, next_cursor = TaskDAO.get_statuses_changed_since( + cursor, task_type=request.args.get("task_type") + ) + return self.response( + 200, + statuses=statuses, + cursor=next_cursor.isoformat() if next_cursor else None, + ) + @expose("//cancel", methods=("POST",)) @protect() @safe diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py index 14b10bfdcf2..bc6abf81303 100644 --- a/superset/tasks/async_queries.py +++ b/superset/tasks/async_queries.py @@ -16,126 +16,187 @@ # under the License. from __future__ import annotations -import dataclasses import logging from typing import Any, TYPE_CHECKING -from celery.exceptions import SoftTimeLimitExceeded from flask import current_app from flask_appbuilder.security.sqla.models import User -from marshmallow import ValidationError +from superset_core.tasks.types import TaskOptions, TaskScope -from superset.charts.data.form_data import set_form_data -from superset.charts.schemas import ChartDataQueryContextSchema -from superset.exceptions import ( - SupersetErrorException, - SupersetErrorsException, +from superset.common.query_serialization import ( + load_serialized_query, + serialize_query, + SerializedQuery, ) +from superset.constants import CacheRegion +from superset.exceptions import SupersetException from superset.extensions import ( - async_query_manager, - celery_app, security_manager, ) +from superset.tasks.decorators import task from superset.utils.core import override_user -from superset.utils.error_sanitization import sanitize_error_dicts if TYPE_CHECKING: from superset.common.query_context import QueryContext + from superset.common.query_object import QueryObject + from superset.models.tasks import Task + from superset.security.guest_token import GuestToken logger = logging.getLogger(__name__) query_timeout = current_app.config[ "SQLLAB_ASYNC_TIME_LIMIT_SEC" ] # TODO: new config key +# GTF task type for the chart-data fan-out. Each QueryObject runs as its own SHARED +# task keyed by its query_cache_key (safe cross-user dedup — the key encodes +# RLS/impersonation). The client polls /api/v1/task/status_changes (filtered to this +# type) and aggregates the tasks' statuses itself; GTF owns completion emission (there +# is no coordinator task). The atomic unit is a QueryObject, not chart-specific, so the +# type is versioned to allow the serialization/execution contract to evolve. +CHART_QUERY_TASK = "superset.query_object_v1" -def _create_query_context_from_form(form_data: dict[str, Any]) -> QueryContext: + +def _resolve_user(user_id: int | None, guest_token: "GuestToken | None") -> User: + """Resolve the acting user for an async chart-data task. + + The GTF executor does not impersonate on its own, so each task establishes the + request user itself (for RLS/impersonation), mirroring the legacy Celery path. """ - Create the query context from the form data. - - :param form_data: The task form data - :returns: The query context - :raises ValidationError: If the request is incorrect - """ - - try: - return ChartDataQueryContextSchema().load(form_data) - except KeyError as ex: - raise ValidationError("Request is incorrect") from ex + if user_id: + return security_manager.get_user_by_id(user_id) + if guest_token: + return security_manager.get_guest_user_from_token(guest_token) + return security_manager.get_anonymous_user() -def _load_user_from_job_metadata(job_metadata: dict[str, Any]) -> User: - if user_id := job_metadata.get("user_id"): - # logged in user - user = security_manager.get_user_by_id(user_id) - elif guest_token := job_metadata.get("guest_token"): - # embedded guest user - user = security_manager.get_guest_user_from_token(guest_token) - del job_metadata["guest_token"] - else: - # default to anonymous user if no user is found - user = security_manager.get_anonymous_user() - return user - - -def _handle_soft_time_limit( - job_metadata: dict[str, Any], ex: Exception, activity: str +def _inject_contribution_totals( + query_obj: "QueryObject", totals_cache_key: str ) -> None: + """Inject ``contribution_totals`` from the cached totals query into ``query_obj``. + + A contribution query normalizes its metrics against column sums from a separate + "totals" query. In the per-query task model the totals query runs as its own + task (a ``depends_on`` prerequisite) and caches its dataframe; here we read that + cached dataframe and inject the sums into this query's contribution + post-processing before it runs — the same result the synchronous + ``ensure_totals_available`` produces, but reading the cache the prerequisite + populated instead of re-running the totals query. ``contribution_totals`` is + stripped from the cache key, so this affects only the result, not the key. """ - SoftTimeLimitExceeded is raised both by a genuine timeout and by a - user-initiated cancel (revoke sends SIGUSR1). The cancel endpoint has - already emitted the terminal event for the latter - it has to, since a task - revoked while still queued never reaches this handler - so only a timeout - is reported here, and without one the client would wait forever. - """ - if async_query_manager.is_job_cancelled(job_metadata["job_id"]): - logger.info("Cancelled by the user while %s", activity) - return + from superset.common.utils.query_cache_manager import QueryCacheManager - logger.warning("A timeout occurred while %s, error: %s", activity, ex) - async_query_manager.update_job( - job_metadata, - async_query_manager.STATUS_ERROR, - errors=[{"message": f"A timeout occurred while {activity}"}], - ) + cache = QueryCacheManager.get(key=totals_cache_key, region=CacheRegion.DATA) + if not cache.is_loaded or cache.df is None: + # The depends_on prerequisite guarantees the totals task succeeded and wrote + # this cache entry, so a miss is unexpected (e.g. it was evicted between the + # totals task finishing and this task reading). Fail loudly rather than + # caching a silently un-normalized result the client would then re-request: + # this task's single query cannot reproduce the synchronous path's + # ensure_totals_available (it has no totals query to run). + raise SupersetException( + f"Contribution totals not found in cache under {totals_cache_key}" + ) + df = cache.df + totals = {col: df[col].sum() for col in df.columns if df[col].dtype.kind in "biufc"} + for post_processing in query_obj.post_processing or []: + if post_processing.get("operation") == "contribution": + post_processing.setdefault("options", {})["contribution_totals"] = totals -@celery_app.task(name="load_chart_data_into_cache", soft_time_limit=query_timeout) -def load_chart_data_into_cache( - job_metadata: dict[str, Any], - form_data: dict[str, Any], +@task(name=CHART_QUERY_TASK, scope=TaskScope.SHARED, timeout=query_timeout) +def execute_chart_query( + serialized_query: SerializedQuery, + user_id: int | None = None, + guest_token: "GuestToken | None" = None, + totals_cache_key: str | None = None, ) -> None: - # pylint: disable=import-outside-toplevel - from superset.commands.chart.data.get_data_command import ChartDataCommand + """Execute a single chart-data query and cache it under its query_cache_key. - with override_user(_load_user_from_job_metadata(job_metadata), force=False): - try: - set_form_data(form_data) - query_context = _create_query_context_from_form(form_data) - command = ChartDataCommand(query_context) - result = command.run(cache=True) - cache_key = result["cache_key"] - result_url = f"/api/v1/chart/data/{cache_key}" - async_query_manager.update_job( - job_metadata, - async_query_manager.STATUS_DONE, - result_url=result_url, + The atomic async unit: reconstruct the one query (canonical serialization), + optionally inject contribution totals from a prerequisite totals task, then run + the existing per-query execution/caching path so a re-request reads the same + DATA-cache entry. + """ + with override_user(_resolve_user(user_id, guest_token), force=False): + query_context = load_serialized_query(serialized_query) + query_obj = query_context.queries[0] + if totals_cache_key: + _inject_contribution_totals(query_obj, totals_cache_key) + # Executes on cache miss and writes CacheRegion.DATA under query_cache_key. + query_context.get_df_payload_result(query_obj) + + +def _query_task_cache_key(query_context: "QueryContext", index: int) -> str | None: + """Compute a query's cache key exactly as its task will. + + ``execute_chart_query`` validates each query before keying (see + ``get_df_payload_result``), so validate here too — otherwise the SHARED task's + ``task_key`` (used for cross-user dedup) could diverge from the key the task + actually caches under. + """ + query_obj = query_context.queries[index] + query_obj.validate() + return query_context.query_cache_key(query_obj) + + +def submit_chart_data_query_tasks( + query_context: "QueryContext", + user_id: int | None, +) -> dict[str, Any]: + """Fan a chart-data request out into one GTF task per ``QueryObject``. + + Each ``QueryObject`` runs as its own SHARED task keyed by its ``query_cache_key`` + (safe cross-user dedup — the key encodes RLS/impersonation), writing the per-query + DATA cache a later re-request reads back. A contribution query ``depends_on`` the + totals query's task and reads its cached result to normalize. + + There is no coordinator task: the client polls ``/api/v1/task/status_changes`` and + aggregates the query tasks' own honest statuses itself (all ``SUCCESS`` → re-issue + the request, now served entirely from the per-query cache; any terminal non-success + → error). GTF owns completion emission (per-task, via the coordination service), so + the websocket transport subscribes to GTF, not to any GAQ-specific stream. + + Returns the HTTP 202 body ``{"task_ids": [...]}`` — the query tasks' UUIDs, in + query order, for the client to poll and cancel via the GTF task API. + """ + guest_user = security_manager.get_current_guest_user_if_guest() + guest_token = guest_user.guest_token if guest_user else None + + queries = query_context.queries + # Contribution queries normalize against a shared totals row. Identify the coupling + # (this also clears the totals query's row_limit so its cache key matches the entry + # its dependents read) and compute the totals key up front. + needs_totals, totals_idx = query_context.prepare_contribution_totals() + totals_key: str | None = None + if needs_totals and totals_idx is not None: + # Mirror the row_limit normalization into the raw serialized dict so the totals + # task caches under the same key its dependents (and the re-request) compute. + query_context.cache_values["queries"][totals_idx]["row_limit"] = None + totals_key = _query_task_cache_key(query_context, totals_idx) + + def _schedule(index: int, depends_on: list["Task"] | None = None) -> "Task": + return execute_chart_query.schedule( + serialize_query(query_context, index), + user_id, + guest_token, + totals_key if index in needs_totals else None, + options=TaskOptions( + task_key=_query_task_cache_key(query_context, index), + depends_on=depends_on, + ), + ) + + # Schedule the totals query first so contribution queries can depend on it. + tasks: dict[int, "Task"] = {} + if totals_idx is not None and needs_totals: + tasks[totals_idx] = _schedule(totals_idx) + for index in range(len(queries)): + if index not in tasks: + depends_on = ( + [tasks[totals_idx]] + if index in needs_totals and totals_idx is not None + else None ) - except SoftTimeLimitExceeded as ex: - _handle_soft_time_limit(job_metadata, ex, "loading chart data") - raise - except Exception as ex: - # Extract SIP-40 style errors when available - if isinstance(ex, SupersetErrorException): - errors = [dataclasses.asdict(ex.error)] - elif isinstance(ex, SupersetErrorsException): - errors = [dataclasses.asdict(error) for error in ex.errors] - else: - # Fallback for non-Superset exceptions - error = str(ex.message if hasattr(ex, "message") else ex) - errors = [{"message": error}] - async_query_manager.update_job( - job_metadata, - async_query_manager.STATUS_ERROR, - errors=sanitize_error_dicts(errors), - ) - raise + tasks[index] = _schedule(index, depends_on=depends_on) + + return {"task_ids": [str(tasks[index].uuid) for index in range(len(queries))]} diff --git a/superset/tasks/decorators.py b/superset/tasks/decorators.py index d8a0e5806ea..6cce2e4efb0 100644 --- a/superset/tasks/decorators.py +++ b/superset/tasks/decorators.py @@ -20,7 +20,16 @@ from __future__ import annotations import inspect import logging -from typing import Any, Callable, cast, Generic, ParamSpec, TYPE_CHECKING, TypeVar +from typing import ( + Any, + Callable, + cast, + Generic, + overload, + ParamSpec, + TYPE_CHECKING, + TypeVar, +) from superset_core.tasks.types import TaskOptions, TaskScope, TaskStatus @@ -42,6 +51,20 @@ P = ParamSpec("P") R = TypeVar("R") +@overload +def task(func: Callable[P, R]) -> "TaskWrapper[P]": ... + + +@overload +def task( + func: None = None, + *, + name: str | None = None, + scope: TaskScope = TaskScope.PRIVATE, + timeout: int | None = None, +) -> Callable[[Callable[P, R]], "TaskWrapper[P]"]: ... + + def task( func: Callable[P, R] | None = None, *, @@ -566,12 +589,21 @@ class TaskWrapper(Generic[P]): if final_task and final_task.status in TERMINAL_STATES: TaskManager.publish_completion(task_uuid, final_task.status) - def schedule(self, *args: P.args, **kwargs: P.kwargs) -> "Task": + def schedule( + self, + *args: Any, + options: TaskOptions | None = None, + **kwargs: Any, + ) -> "Task": """ Schedule this task for asynchronous execution. The signature mirrors the original task function, with an additional - keyword-only 'options' parameter for execution metadata. + keyword-only 'options' parameter for execution metadata. Business args + are typed as ``Any`` here because PEP 612 does not allow adding a + keyword-only parameter alongside the captured ``ParamSpec``; the runtime + ``__signature__`` (patched in ``__init__``) still mirrors the wrapped + function for introspection and IDE support. Args: *args, **kwargs: Business arguments for the task function @@ -607,7 +639,7 @@ class TaskWrapper(Generic[P]): raise GlobalTaskFrameworkDisabledError() # Extract and merge options (decorator defaults + call-time overrides) - override_options = cast(TaskOptions | None, kwargs.pop("options", None)) + override_options = options options = self._merge_options(override_options) # Validate task configuration diff --git a/superset/tasks/filters.py b/superset/tasks/filters.py index 9159a465969..066bffdecb8 100644 --- a/superset/tasks/filters.py +++ b/superset/tasks/filters.py @@ -52,11 +52,27 @@ class TaskFilter(BaseFilter): # pylint: disable=too-few-public-methods from superset import security_manager from superset.models.task_subscribers import TaskSubscriber from superset.models.tasks import Task + from superset.tasks.guest import get_current_guest_subscriber_key user_id = get_user_id() if not user_id: - # Within a request, a principal without a user id gets no tasks; - # background jobs run outside a request context and are unfiltered. + # Embedded guests have no ab_user id but subscribe by a token-derived + # key, so scope their visibility to tasks carrying that key. + if guest_key := get_current_guest_subscriber_key(): + guest_subscribed = ( + select(TaskSubscriber.id) + .where( + and_( + TaskSubscriber.task_id == Task.id, + TaskSubscriber.guest_key == guest_key, + ) + ) + .exists() + ) + return query.filter(guest_subscribed) + # A principal without a user id or guest identity gets no tasks + # within a request; background jobs run outside a request and are + # unfiltered. if has_request_context(): return query.filter(false()) return query diff --git a/superset/tasks/guest.py b/superset/tasks/guest.py new file mode 100644 index 00000000000..56bb433ffb0 --- /dev/null +++ b/superset/tasks/guest.py @@ -0,0 +1,72 @@ +# 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. +"""Guest identity for Global Task Framework subscriptions. + +Embedded guest users have no ``ab_user`` row, so they cannot subscribe to tasks +by ``user_id``. Instead they subscribe by a ``guest_key``: a stable, unguessable +identity derived from their guest token, which the task filter honors to grant a +guest visibility of the tasks it created or (via SHARED-scope dedup) joined. +""" + +from __future__ import annotations + +import hashlib +import hmac + +from flask import current_app + +from superset import security_manager +from superset.utils import json + + +def get_current_guest_subscriber_key() -> str | None: + """Return a stable subscriber key for the current guest, or ``None``. + + ``None`` when the request is not an embedded guest (an authenticated user + subscribes by ``user_id`` instead). The key is an HMAC over the guest token's + stable identifying claims, keyed with the app ``SECRET_KEY`` so it is + unguessable to outside callers and reproducible for the same token across the + request that schedules a task and the polls that await it. + """ + guest_user = security_manager.get_current_guest_user_if_guest() + if not guest_user: + return None + token = guest_user.guest_token + # Bind the key to every authorization-relevant claim so two tokens that differ + # in their effective access scope derive different keys (and can't see each + # other's tasks): ``iat``/``exp`` pin it to a single issuance, ``resources``/ + # ``datasets``/``rev`` to the granted resources, and ``rls_rules`` to the + # row-level scope. + message = json.dumps( + { + "user": token.get("user"), + "resources": token.get("resources"), + "iat": token.get("iat"), + "exp": token.get("exp"), + "aud": token.get("aud"), + "datasets": token.get("datasets"), + "rev": token.get("rev"), + "rls_rules": token.get("rls_rules"), + }, + sort_keys=True, + ).encode("utf-8") + digest = hmac.new( + current_app.config["SECRET_KEY"].encode("utf-8"), + message, + hashlib.sha256, + ).hexdigest() + return f"guest-{digest}" diff --git a/superset/views/base.py b/superset/views/base.py index f938b32b33c..9fc70acda66 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -90,12 +90,10 @@ FRONTEND_CONF_KEYS = ( "SQLLAB_SAVE_WARNING_MESSAGE", "SQLLAB_DEFAULT_DBID", "DISPLAY_MAX_ROW", - "GLOBAL_ASYNC_QUERIES_TRANSPORT", "GLOBAL_ASYNC_QUERIES_POLLING_DELAY", "SQL_VALIDATORS_BY_ENGINE", "SQLALCHEMY_DOCS_URL", "SQLALCHEMY_DISPLAY_TEXT", - "GLOBAL_ASYNC_QUERIES_WEBSOCKET_URL", "DASHBOARD_AUTO_REFRESH_MODE", "DASHBOARD_AUTO_REFRESH_INTERVALS", "DASHBOARD_VIRTUALIZATION", diff --git a/tests/integration_tests/async_events/api_tests.py b/tests/integration_tests/async_events/api_tests.py deleted file mode 100644 index 89e008940bd..00000000000 --- a/tests/integration_tests/async_events/api_tests.py +++ /dev/null @@ -1,214 +0,0 @@ -# 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. -from typing import Any, Optional, Type -from unittest import mock - -from superset.async_events.async_query_manager import ( - AsyncQueryJobException, - AsyncQueryTokenException, -) -from superset.async_events.cache_backend import ( - RedisCacheBackend, - RedisSentinelCacheBackend, -) -from superset.extensions import async_query_manager, async_query_manager_factory -from superset.utils import json -from tests.integration_tests.base_tests import SupersetTestCase -from tests.integration_tests.constants import ADMIN_USERNAME -from tests.integration_tests.test_app import app - - -class TestAsyncEventApi(SupersetTestCase): - UUID = "943c920-32a5-412a-977d-b8e47d36f5a4" - JOB_ID = "10a0bd9a-03c8-4737-9345-f4234ba86512" - - def fetch_events(self, last_id: Optional[str] = None): - base_uri = "api/v1/async_event/" - uri = f"{base_uri}?last_id={last_id}" if last_id else base_uri - return self.client.get(uri) - - def cancel_event(self, job_id: str): - return self.client.post(f"api/v1/async_event/{job_id}/cancel") - - def run_test_with_cache_backend(self, cache_backend_cls: Type[Any], test_func): - app._got_first_request = False - - # GAQ resolves its own backend from get_cache_backend during init_app, so - # inject the mock there. init_app must run before login(): it re-registers the - # after_request handler, and login()'s request would otherwise trip Flask's - # "setup method after first request" guard before init_app gets to re-register. - mock_cache = mock.Mock(spec=cache_backend_cls) - with mock.patch( - "superset.async_events.async_query_manager.get_cache_backend", - return_value=mock_cache, - ): - async_query_manager_factory.init_app(app) - - self.login(ADMIN_USERNAME) - test_func(mock_cache) - - def _test_events_logic(self, mock_cache): - with mock.patch.object(mock_cache, "xrange") as mock_xrange: - rv = self.fetch_events() - response = json.loads(rv.data.decode("utf-8")) - - assert rv.status_code == 200 - channel_id = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] + self.UUID - mock_xrange.assert_called_with(channel_id, "-", "+", 100) - assert response == {"result": []} - - def _test_events_last_id_logic(self, mock_cache): - with mock.patch.object(mock_cache, "xrange") as mock_xrange: - rv = self.fetch_events("1607471525180-0") - response = json.loads(rv.data.decode("utf-8")) - - assert rv.status_code == 200 - channel_id = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] + self.UUID - mock_xrange.assert_called_with(channel_id, "1607471525180-1", "+", 100) - assert response == {"result": []} - - def _test_events_results_logic(self, mock_cache): - with mock.patch.object(mock_cache, "xrange") as mock_xrange: - mock_xrange.return_value = [ - ( - "1607477697866-0", - { - "data": '{"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", "job_id": "10a0bd9a-03c8-4737-9345-f4234ba86512", "user_id": "1", "status": "done", "errors": [], "result_url": "/api/v1/chart/data/qc-ecd766dd461f294e1bcdaa321e0e8463"}' # noqa: E501 - }, - ), - ( - "1607477697993-0", - { - "data": '{"channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", "job_id": "027cbe49-26ce-4813-bb5a-0b95a626b84c", "user_id": "1", "status": "done", "errors": [], "result_url": "/api/v1/chart/data/qc-1bbc3a240e7039ba4791aefb3a7ee80d"}' # noqa: E501 - }, - ), - ] - rv = self.fetch_events() - response = json.loads(rv.data.decode("utf-8")) - - assert rv.status_code == 200 - channel_id = app.config["GLOBAL_ASYNC_QUERIES_REDIS_STREAM_PREFIX"] + self.UUID - mock_xrange.assert_called_with(channel_id, "-", "+", 100) - expected = { - "result": [ - { - "channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", - "errors": [], - "id": "1607477697866-0", - "job_id": "10a0bd9a-03c8-4737-9345-f4234ba86512", - "result_url": "/api/v1/chart/data/qc-ecd766dd461f294e1bcdaa321e0e8463", # noqa: E501 - "status": "done", - "user_id": "1", - }, - { - "channel_id": "1095c1c9-b6b1-444d-aa83-8e323b32831f", - "errors": [], - "id": "1607477697993-0", - "job_id": "027cbe49-26ce-4813-bb5a-0b95a626b84c", - "result_url": "/api/v1/chart/data/qc-1bbc3a240e7039ba4791aefb3a7ee80d", # noqa: E501 - "status": "done", - "user_id": "1", - }, - ] - } - assert response == expected - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_events_redis_cache_backend(self, mock_uuid4): - self.run_test_with_cache_backend(RedisCacheBackend, self._test_events_logic) - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_events_redis_sentinel_cache_backend(self, mock_uuid4): - self.run_test_with_cache_backend( - RedisSentinelCacheBackend, self._test_events_logic - ) - - def test_events_no_login(self): - app._got_first_request = False - async_query_manager_factory.init_app(app) - rv = self.fetch_events() - assert rv.status_code == 401 - - def test_events_no_token(self): - self.login(ADMIN_USERNAME) - self.client.set_cookie(app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"], "") - rv = self.fetch_events() - assert rv.status_code == 401 - - def _test_cancel_logic(self, mock_cache): - with mock.patch.object(async_query_manager, "cancel_job") as mock_cancel: - rv = self.cancel_event(self.JOB_ID) - - assert rv.status_code == 200 - mock_cancel.assert_called_once() - assert mock_cancel.call_args.args[0] == self.JOB_ID - response = json.loads(rv.data.decode("utf-8")) - assert response["result"] == {"job_id": self.JOB_ID, "status": "cancelled"} - - def _test_cancel_invalid_job_id_logic(self, mock_cache): - # job_id lands in a Redis key, so reject anything that isn't a UUID - # before it reaches the cache backend. - with mock.patch.object(async_query_manager, "cancel_job") as mock_cancel: - rv = self.cancel_event("not-a-uuid") - assert rv.status_code == 400 - mock_cancel.assert_not_called() - - def _test_cancel_forbidden_logic(self, mock_cache): - with mock.patch.object( - async_query_manager, - "cancel_job", - side_effect=AsyncQueryTokenException("nope"), - ): - rv = self.cancel_event(self.JOB_ID) - assert rv.status_code == 403 - - def _test_cancel_not_found_logic(self, mock_cache): - with mock.patch.object( - async_query_manager, - "cancel_job", - side_effect=AsyncQueryJobException("gone"), - ): - rv = self.cancel_event(self.JOB_ID) - assert rv.status_code == 404 - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_cancel_redis_cache_backend(self, mock_uuid4): - self.run_test_with_cache_backend(RedisCacheBackend, self._test_cancel_logic) - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_cancel_invalid_job_id(self, mock_uuid4): - self.run_test_with_cache_backend( - RedisCacheBackend, self._test_cancel_invalid_job_id_logic - ) - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_cancel_forbidden(self, mock_uuid4): - self.run_test_with_cache_backend( - RedisCacheBackend, self._test_cancel_forbidden_logic - ) - - @mock.patch("uuid.uuid4", return_value=UUID) - def test_cancel_not_found(self, mock_uuid4): - self.run_test_with_cache_backend( - RedisCacheBackend, self._test_cancel_not_found_logic - ) - - def test_cancel_no_login(self): - app._got_first_request = False - async_query_manager_factory.init_app(app) - rv = self.cancel_event(self.JOB_ID) - assert rv.status_code == 401 diff --git a/tests/integration_tests/charts/data/api_tests.py b/tests/integration_tests/charts/data/api_tests.py index 18a30c543fd..282999a6af8 100644 --- a/tests/integration_tests/charts/data/api_tests.py +++ b/tests/integration_tests/charts/data/api_tests.py @@ -45,7 +45,7 @@ from superset.common.chart_data_timing import ( from superset.connectors.sqla.models import SqlaTable, TableColumn from superset.constants import CACHE_DISABLED_TIMEOUT from superset.errors import SupersetErrorType -from superset.extensions import async_query_manager_factory, db +from superset.extensions import db from superset.models.annotations import AnnotationLayer from superset.models.slice import Slice from superset.models.sql_lab import Query @@ -64,7 +64,7 @@ from tests.conftest import with_config from tests.integration_tests.annotation_layers.fixtures import ( create_annotation_layers, # noqa: F401 ) -from tests.integration_tests.base_tests import SupersetTestCase, test_client +from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.constants import ( ADMIN_USERNAME, @@ -762,8 +762,6 @@ class TestPostChartDataApi(BaseTestChartDataApi): @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_chart_data_async(self): self.logout() - app._got_first_request = False - async_query_manager_factory.init_app(app) self.login(ADMIN_USERNAME) # Introducing time.sleep to make test less flaky with MySQL time.sleep(1) @@ -772,10 +770,10 @@ class TestPostChartDataApi(BaseTestChartDataApi): assert rv.status_code == 202 time.sleep(1) data = json.loads(rv.data.decode("utf-8")) - keys = list(data.keys()) - self.assertCountEqual( # noqa: PT009 - keys, ["channel_id", "job_id", "user_id", "status", "errors", "result_url"] - ) + # The async response is the GTF job: the per-QueryObject task uuids the + # client polls via /api/v1/task/status_changes. + assert list(data.keys()) == ["task_ids"] + assert isinstance(data["task_ids"], list) @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") @@ -785,8 +783,6 @@ class TestPostChartDataApi(BaseTestChartDataApi): Chart data API: Test chart data query returns results synchronously when results are already cached, and that is_cached is logged. """ - app._got_first_request = False - async_query_manager_factory.init_app(app) class QueryContext: result_format = ChartDataResultFormat.JSON @@ -869,8 +865,6 @@ class TestPostChartDataApi(BaseTestChartDataApi): """ Chart data API: Test that force=true skips cache and triggers async job """ - app._got_first_request = False - async_query_manager_factory.init_app(app) # Mock the command execution to return cached data class QueryContext: @@ -904,10 +898,7 @@ class TestPostChartDataApi(BaseTestChartDataApi): # since we skip the cache check entirely mock_execute.assert_not_called() data = json.loads(rv.data.decode("utf-8")) - keys = list(data.keys()) - self.assertCountEqual( # noqa: PT009 - keys, ["channel_id", "job_id", "user_id", "status", "errors", "result_url"] - ) + assert list(data.keys()) == ["task_ids"] @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") @@ -915,26 +906,10 @@ class TestPostChartDataApi(BaseTestChartDataApi): """ Chart data API: Test chart data query non-JSON format (async) """ - app._got_first_request = False - async_query_manager_factory.init_app(app) self.query_context_payload["result_type"] = "results" rv = self.post_assert_metric(CHART_DATA_URI, self.query_context_payload, "data") assert rv.status_code == 200 - @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") - def test_chart_data_async_invalid_token(self): - """ - Chart data API: Test chart data query (async) - """ - app._got_first_request = False - async_query_manager_factory.init_app(app) - test_client.set_cookie( - app.config["GLOBAL_ASYNC_QUERIES_JWT_COOKIE_NAME"], "foo" - ) - rv = test_client.post(CHART_DATA_URI, json=self.query_context_payload) - assert rv.status_code == 401 - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_chart_data_rowcount(self): """ @@ -1436,92 +1411,6 @@ class TestGetChartDataApi(BaseTestChartDataApi): # is_cached should be [True] when retrieved from cache assert records[0]["is_cached"] == [True] - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") - @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) - @mock.patch("superset.charts.data.api.QueryContextCacheLoader") - def test_chart_data_cache(self, cache_loader): - """ - Chart data cache API: Test chart data async cache request - """ - app._got_first_request = False - async_query_manager_factory.init_app(app) - cache_loader.load.return_value = self.query_context_payload - orig_execute = ChartDataCommand.execute - - def mock_execute(self, **kwargs): - assert kwargs["force_cached"] is True # noqa: E712 - # override force_cached to get result from DB - return orig_execute(self, force_cached=False) - - with mock.patch.object(ChartDataCommand, "execute", new=mock_execute): - rv = self.get_assert_metric( - f"{CHART_DATA_URI}/test-cache-key", "data_from_cache" - ) - data = json.loads(rv.data.decode("utf-8")) - - expected_row_count = self.get_expected_row_count("client_id_3") - assert rv.status_code == 200 - assert data["result"][0]["rowcount"] == expected_row_count - - @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) - @mock.patch("superset.charts.data.api.QueryContextCacheLoader") - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") - def test_chart_data_cache_run_failed(self, cache_loader): - """ - Chart data cache API: Test chart data async cache request with run failure - """ - app._got_first_request = False - async_query_manager_factory.init_app(app) - cache_loader.load.return_value = self.query_context_payload - rv = self.get_assert_metric( - f"{CHART_DATA_URI}/test-cache-key", "data_from_cache" - ) - data = json.loads(rv.data.decode("utf-8")) - - assert rv.status_code == 422 - assert data["message"] == "Error loading data from cache" - - @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) - @mock.patch("superset.charts.data.api.QueryContextCacheLoader") - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") - def test_chart_data_cache_no_login(self, cache_loader): - """ - Chart data cache API: Test chart data async cache request (no login) - """ - if get_example_database().backend == "presto": - return - - app._got_first_request = False - async_query_manager_factory.init_app(app) - self.logout() - cache_loader.load.return_value = self.query_context_payload - orig_execute = ChartDataCommand.execute - - def mock_execute(self, **kwargs): - assert kwargs["force_cached"] is True # noqa: E712 - # override force_cached to get result from DB - return orig_execute(self, force_cached=False) - - with mock.patch.object(ChartDataCommand, "execute", new=mock_execute): - rv = self.client.get( - f"{CHART_DATA_URI}/test-cache-key", - ) - - assert rv.status_code == 401 - - @with_feature_flags(GLOBAL_ASYNC_QUERIES=True) - def test_chart_data_cache_key_error(self): - """ - Chart data cache API: Test chart data async cache request with invalid cache key - """ - app._got_first_request = False - async_query_manager_factory.init_app(app) - rv = self.get_assert_metric( - f"{CHART_DATA_URI}/test-cache-key", "data_from_cache" - ) - - assert rv.status_code == 404 - @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_chart_data_with_adhoc_column(self): """ diff --git a/tests/integration_tests/superset_test_config.py b/tests/integration_tests/superset_test_config.py index eb6567a7478..11ba652d392 100644 --- a/tests/integration_tests/superset_test_config.py +++ b/tests/integration_tests/superset_test_config.py @@ -155,7 +155,6 @@ EXPLORE_FORM_DATA_CACHE_CONFIG = { "CACHE_DEFAULT_TIMEOUT": int(timedelta(minutes=10).total_seconds()), } -GLOBAL_ASYNC_QUERIES_JWT_SECRET = "test-secret-change-me-test-secret-change-me" # noqa: S105 ALERT_REPORTS_WORKING_TIME_OUT_KILL = True diff --git a/tests/integration_tests/tasks/async_queries_tests.py b/tests/integration_tests/tasks/async_queries_tests.py deleted file mode 100644 index f254840804d..00000000000 --- a/tests/integration_tests/tasks/async_queries_tests.py +++ /dev/null @@ -1,119 +0,0 @@ -# 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. -"""Unit tests for async query celery jobs in Superset""" - -from unittest import mock -from uuid import uuid4 - -import pytest -from celery.exceptions import SoftTimeLimitExceeded - -from superset.commands.chart.data.get_data_command import ChartDataCommand -from superset.commands.chart.exceptions import ChartDataQueryFailedError -from superset.extensions import async_query_manager, security_manager -from tests.integration_tests.base_tests import SupersetTestCase -from tests.integration_tests.fixtures.birth_names_dashboard import ( - load_birth_names_dashboard_with_slices, # noqa: F401 - load_birth_names_data, # noqa: F401 -) -from tests.integration_tests.fixtures.query_context import get_query_context -from tests.integration_tests.fixtures.tags import ( - with_tagging_system_feature, # noqa: F401 -) -from tests.integration_tests.test_app import app - - -@pytest.mark.usefixtures( - "load_birth_names_data", "load_birth_names_dashboard_with_slices" -) -class TestAsyncQueries(SupersetTestCase): - @mock.patch("superset.tasks.async_queries.set_form_data") - @mock.patch.object(async_query_manager, "update_job") - def test_load_chart_data_into_cache(self, mock_update_job, mock_set_form_data): - from superset.tasks.async_queries import load_chart_data_into_cache - - app._got_first_request = False - - query_context = get_query_context("birth_names") - user = security_manager.find_user("gamma") - job_metadata = { - "channel_id": str(uuid4()), - "job_id": str(uuid4()), - "user_id": user.id, - "status": "pending", - "errors": [], - } - - load_chart_data_into_cache(job_metadata, query_context) - - mock_set_form_data.assert_called_once_with(query_context) - mock_update_job.assert_called_once_with( - job_metadata, "done", result_url=mock.ANY - ) - - @mock.patch.object( - ChartDataCommand, "run", side_effect=ChartDataQueryFailedError("Error: foo") - ) - @mock.patch.object(async_query_manager, "update_job") - def test_load_chart_data_into_cache_error(self, mock_update_job, mock_run_command): - from superset.tasks.async_queries import load_chart_data_into_cache - - app._got_first_request = False - - query_context = get_query_context("birth_names") - user = security_manager.find_user("gamma") - job_metadata = { - "channel_id": str(uuid4()), - "job_id": str(uuid4()), - "user_id": user.id, - "status": "pending", - "errors": [], - } - with pytest.raises(ChartDataQueryFailedError): - load_chart_data_into_cache(job_metadata, query_context) - - mock_run_command.assert_called_once_with(cache=True) - errors = [{"message": "Error: foo"}] - mock_update_job.assert_called_once_with(job_metadata, "error", errors=errors) - - @mock.patch.object(ChartDataCommand, "run") - @mock.patch.object(async_query_manager, "update_job") - def test_soft_timeout_load_chart_data_into_cache( - self, mock_update_job, mock_run_command - ): - from superset.tasks.async_queries import load_chart_data_into_cache - - app._got_first_request = False - - user = security_manager.find_user("gamma") - form_data = {} - job_metadata = { - "channel_id": str(uuid4()), - "job_id": str(uuid4()), - "user_id": user.id, - "status": "pending", - "errors": [], - } - errors = ["A timeout occurred while loading chart data"] - - with pytest.raises(SoftTimeLimitExceeded): # noqa: PT012 - with mock.patch( - "superset.tasks.async_queries.set_form_data" - ) as set_form_data: - set_form_data.side_effect = SoftTimeLimitExceeded() - load_chart_data_into_cache(job_metadata, form_data) - set_form_data.assert_called_once_with(form_data, "error", errors=errors) diff --git a/tests/unit_tests/async_events/async_query_manager_tests.py b/tests/unit_tests/async_events/async_query_manager_tests.py deleted file mode 100644 index 91dbd267e05..00000000000 --- a/tests/unit_tests/async_events/async_query_manager_tests.py +++ /dev/null @@ -1,538 +0,0 @@ -# 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. -from datetime import datetime, timedelta, timezone -from unittest import mock -from unittest.mock import ANY, Mock - -from flask import g -from jwt import encode -from pytest import fixture, raises # noqa: PT013 - -from superset import security_manager -from superset.async_events.async_query_manager import ( - AsyncQueryJobException, - AsyncQueryManager, - AsyncQueryTokenException, -) -from superset.async_events.cache_backend import ( - RedisCacheBackend, -) -from superset.utils import json - -JWT_TOKEN_SECRET = "some_secret" # noqa: S105 -JWT_TOKEN_COOKIE_NAME = "superset_async_jwt" # noqa: S105 - - -@fixture -def async_query_manager(): - query_manager = AsyncQueryManager() - query_manager._jwt_secret = JWT_TOKEN_SECRET - query_manager._jwt_cookie_name = JWT_TOKEN_COOKIE_NAME - query_manager._jwt_expiration_seconds = 3600 - return query_manager - - -def set_current_as_guest_user(): - g.user = security_manager.get_guest_user_from_token( - { - "user": {}, - "resources": [{"type": "dashboard", "id": "some-uuid"}], - "rls_rules": [{"clause": '"STATEID" = 3'}], - "iat": 1700000000.0, - "exp": 1700000300.0, - "aud": "http://0.0.0.0:8080/", - "type": "guest", - } - ) - - -def test_parse_channel_id_from_request(async_query_manager): - encoded_token = encode( - {"channel": "test_channel_id"}, JWT_TOKEN_SECRET, algorithm="HS256" - ) - - request = Mock() - request.cookies = {"superset_async_jwt": encoded_token} - - assert ( - async_query_manager.parse_channel_id_from_request(request) == "test_channel_id" - ) - - -def test_parse_channel_id_from_request_with_valid_exp(async_query_manager): - """A token with a future exp claim is accepted.""" - encoded_token = encode( - { - "channel": "test_channel_id", - "exp": datetime.now(tz=timezone.utc) + timedelta(hours=1), - }, - JWT_TOKEN_SECRET, - algorithm="HS256", - ) - - request = Mock() - request.cookies = {"superset_async_jwt": encoded_token} - - assert ( - async_query_manager.parse_channel_id_from_request(request) == "test_channel_id" - ) - - -def test_parse_channel_id_from_request_expired_token(async_query_manager): - """A token with a past exp claim is rejected by the decode path.""" - encoded_token = encode( - { - "channel": "test_channel_id", - "exp": datetime.now(tz=timezone.utc) - timedelta(seconds=1), - }, - JWT_TOKEN_SECRET, - algorithm="HS256", - ) - - request = Mock() - request.cookies = {"superset_async_jwt": encoded_token} - - with raises(AsyncQueryTokenException): - async_query_manager.parse_channel_id_from_request(request) - - -def test_init_app_issues_token_with_exp_claim(): - """Tokens issued through the request handler carry an exp claim.""" - import jwt - - app = Mock() - app.config = { - "GLOBAL_ASYNC_QUERIES_JWT_SECRET": JWT_TOKEN_SECRET, - "GLOBAL_ASYNC_QUERIES_JWT_EXPIRATION_SECONDS": 3600, - } - query_manager = AsyncQueryManager() - query_manager._jwt_secret = app.config["GLOBAL_ASYNC_QUERIES_JWT_SECRET"] - query_manager._jwt_expiration_seconds = app.config[ - "GLOBAL_ASYNC_QUERIES_JWT_EXPIRATION_SECONDS" - ] - - before = datetime.now(tz=timezone.utc) - token = encode( - { - "channel": "test_channel_id", - "exp": before + timedelta(seconds=query_manager._jwt_expiration_seconds), - }, - query_manager._jwt_secret, - algorithm="HS256", - ) - decoded = jwt.decode(token, JWT_TOKEN_SECRET, algorithms=["HS256"]) - assert "exp" in decoded - assert decoded["exp"] >= int(before.timestamp()) - - -def test_parse_channel_id_from_request_no_cookie(async_query_manager): - request = Mock() - request.cookies = {} - - with raises(AsyncQueryTokenException): - async_query_manager.parse_channel_id_from_request(request) - - -def test_parse_channel_id_from_request_bad_jwt(async_query_manager): - request = Mock() - request.cookies = {"superset_async_jwt": "bad_jwt"} - - with raises(AsyncQueryTokenException): - async_query_manager.parse_channel_id_from_request(request) - - -@mock.patch("superset.is_feature_enabled") -def test_parse_channel_id_from_request_as_guest_user_no_cookie( - is_feature_enabled_mock, async_query_manager -): - """ - Embedded guest sessions cannot rely on the async-token cookie because - cross-origin cookies are blocked or stripped by modern browsers when the - dashboard is rendered inside a third-party iframe. The channel id must - therefore be derived from the guest token rather than the cookie. - """ - is_feature_enabled_mock.return_value = True - set_current_as_guest_user() - - request = Mock() - request.cookies = {} - - channel_id = async_query_manager.parse_channel_id_from_request(request) - assert channel_id.startswith("guest-") - - -@mock.patch("superset.is_feature_enabled") -def test_parse_channel_id_from_request_as_guest_user_is_deterministic( - is_feature_enabled_mock, async_query_manager -): - """ - The same guest token (including its RLS rules) must yield the same channel - id across requests. Otherwise the chart-data submission and the polling - endpoint would write to and read from different streams, returning 401s - even though the work was scheduled correctly. - """ - is_feature_enabled_mock.return_value = True - set_current_as_guest_user() - - request = Mock() - request.cookies = {} - - first = async_query_manager.parse_channel_id_from_request(request) - second = async_query_manager.parse_channel_id_from_request(request) - assert first == second - - -@mock.patch("superset.is_feature_enabled") -def test_parse_channel_id_from_request_as_guest_user_differs_per_token( - is_feature_enabled_mock, async_query_manager -): - """Different guest tokens must produce different channel ids.""" - is_feature_enabled_mock.return_value = True - - set_current_as_guest_user() - request = Mock() - request.cookies = {} - first = async_query_manager.parse_channel_id_from_request(request) - - g.user = security_manager.get_guest_user_from_token( - { - "user": {"username": "other"}, - "resources": [{"type": "dashboard", "id": "another-uuid"}], - "rls_rules": [{"clause": '"STATEID" = 4'}], - "iat": 1700000000.0, - "exp": 1700000300.0, - "aud": "http://0.0.0.0:8080/", - "type": "guest", - } - ) - second = async_query_manager.parse_channel_id_from_request(request) - - assert first != second - - -@mock.patch("superset.is_feature_enabled") -def test_parse_channel_id_from_request_as_guest_user_differs_per_scope( - is_feature_enabled_mock, async_query_manager -): - """ - Tokens that differ only in the optional ``datasets`` allowlist or ``rev`` - revocation version must still derive distinct channel ids, otherwise two - differently scoped embedded sessions would collide on the same stream. - """ - is_feature_enabled_mock.return_value = True - - base_token = { - "user": {}, - "resources": [{"type": "dashboard", "id": "some-uuid"}], - "rls_rules": [{"clause": '"STATEID" = 3'}], - "iat": 1700000000.0, - "exp": 1700000300.0, - "aud": "http://0.0.0.0:8080/", - "type": "guest", - } - - request = Mock() - request.cookies = {} - - g.user = security_manager.get_guest_user_from_token(dict(base_token)) - baseline = async_query_manager.parse_channel_id_from_request(request) - - g.user = security_manager.get_guest_user_from_token({**base_token, "datasets": [1]}) - with_datasets = async_query_manager.parse_channel_id_from_request(request) - - g.user = security_manager.get_guest_user_from_token({**base_token, "rev": 1}) - with_rev = async_query_manager.parse_channel_id_from_request(request) - - assert baseline != with_datasets - assert baseline != with_rev - assert with_datasets != with_rev - - -@mock.patch("superset.is_feature_enabled") -def test_submit_chart_data_job_as_guest_user( - is_feature_enabled_mock, async_query_manager -): - is_feature_enabled_mock.return_value = True - set_current_as_guest_user() - - # The manager has no GAQ backend wired (init_app not run), so the best-effort - # cancel registry write is skipped and submission still proceeds. - job_mock = Mock() - async_query_manager._load_chart_data_into_cache_job = job_mock - job_meta = async_query_manager.submit_chart_data_job( - channel_id="test_channel_id", - form_data={}, - ) - - job_mock.apply_async.assert_called_once_with( - args=[ - { - "channel_id": "test_channel_id", - "errors": [], - "guest_token": { - "user": {}, - "resources": [{"type": "dashboard", "id": "some-uuid"}], - "rls_rules": [{"clause": '"STATEID" = 3'}], - "iat": 1700000000.0, - "exp": 1700000300.0, - "aud": "http://0.0.0.0:8080/", - "type": "guest", - }, - "job_id": ANY, - "result_url": None, - "status": "pending", - "user_id": None, - }, - {}, - ], - task_id=ANY, - expires=3600, - ) - - assert "guest_token" not in job_meta - job_mock.reset_mock() # Reset the mock for the next iteration - - -def test_parse_channel_id_from_request_sub_none(async_query_manager): - """Regression: token with sub=None must not break parse (PyJWT 2.10.1+).""" - encoded_token = encode( - {"channel": "test_channel_id", "sub": None}, - JWT_TOKEN_SECRET, - algorithm="HS256", - ) - - request = Mock() - request.cookies = {JWT_TOKEN_COOKIE_NAME: encoded_token} - - with raises(AsyncQueryTokenException): - async_query_manager.parse_channel_id_from_request(request) - - -def test_validate_session_guest_user_creates_valid_token(async_query_manager): - """Regression: validate_session creates decodable tokens when user_id is None.""" - from flask import Flask - - async_query_manager._jwt_cookie_secure = False - async_query_manager._jwt_cookie_domain = None - async_query_manager._jwt_cookie_samesite = "Lax" - async_query_manager._jwt_expiration_seconds = 3600 - - app = Flask(__name__) - app.secret_key = "test_secret_key_for_testing" # noqa: S105 - async_query_manager.register_request_handlers(app) - - @app.route("/test") - def test_view(): - return "ok" - - with mock.patch( - "superset.async_events.async_query_manager.get_user_id", - return_value=None, - ): - client = app.test_client() - resp = client.get("/test") - - cookie_header = [ - v - for k, v in resp.headers - if k == "Set-Cookie" and JWT_TOKEN_COOKIE_NAME in v - ] - assert cookie_header, "JWT cookie was not set" - token = cookie_header[0].split("=", 1)[1].split(";")[0] - - mock_request = Mock() - mock_request.cookies = {JWT_TOKEN_COOKIE_NAME: token} - channel = async_query_manager.parse_channel_id_from_request(mock_request) - assert channel # valid UUID string - - -@fixture -def coordination_backend(): - """A mock Redis backend GAQ operates against as its own dedicated backend.""" - return mock.Mock(spec=RedisCacheBackend) - - -@fixture -def cancellable_manager(coordination_backend): - """A manager wired to a mock Redis backend for cancellation tests.""" - manager = AsyncQueryManager() - manager._jwt_expiration_seconds = 3600 - manager._stream_prefix = "async-events-" - manager._gaq_backend = coordination_backend - return manager - - -def test_init_job_registers_cancellable_record( - cancellable_manager, coordination_backend -): - """init_job persists the owner identity a later cancel must match.""" - cancellable_manager.init_job("chan-1", 7) - - coordination_backend.set.assert_called_once() - key, value = coordination_backend.set.call_args.args - assert key.startswith("async-events-job-cancel:") - assert json.loads(value) == {"channel_id": "chan-1", "user_id": 7} - - -def test_cancel_job_authorized_revokes_task(cancellable_manager, coordination_backend): - cancellable_manager._stream_limit = 100 - cancellable_manager._stream_limit_firehose = 1000 - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - coordination_backend.set.return_value = True - - with mock.patch("superset.extensions.celery_app") as celery_app: - cancellable_manager.cancel_job("job-1", "chan-1", 7) - - celery_app.control.revoke.assert_called_once_with( - "job-1", terminate=True, signal="SIGUSR1" - ) - # The job is flagged cancelled (conditionally, xx=True) so the worker knows - # what the signal it is about to receive means. - assert coordination_backend.set.call_args.kwargs["xx"] is True - flagged = json.loads(coordination_backend.set.call_args.args[1]) - assert flagged["cancelled"] is True - - -def test_cancel_job_emits_the_terminal_event(cancellable_manager, coordination_backend): - """A task revoked before a worker picks it up never reports on itself.""" - cancellable_manager._stream_limit = 100 - cancellable_manager._stream_limit_firehose = 1000 - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - coordination_backend.set.return_value = True - - with mock.patch("superset.extensions.celery_app"): - cancellable_manager.cancel_job("job-1", "chan-1", 7) - - scoped_stream, event_data = coordination_backend.xadd.call_args_list[0].args[:2] - assert scoped_stream == "async-events-chan-1" - assert json.loads(event_data["data"]) == { - "channel_id": "chan-1", - "job_id": "job-1", - "user_id": 7, - "status": AsyncQueryManager.STATUS_CANCELLED, - "errors": [], - "result_url": None, - } - # the record has to outlive the event: the worker still needs to recognize - # the signal on its way as a cancellation - coordination_backend.delete.assert_not_called() - - -def test_cancel_job_completed_between_read_and_flag( - cancellable_manager, coordination_backend -): - """If the job's record is cleared after the auth read, don't revoke.""" - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - # Conditional (xx) write finds no key: the job finished and cleaned up. - coordination_backend.set.return_value = None - - with ( - mock.patch("superset.extensions.celery_app") as celery_app, - raises(AsyncQueryJobException), - ): - cancellable_manager.cancel_job("job-1", "chan-1", 7) - - celery_app.control.revoke.assert_not_called() - - -def test_cancel_job_wrong_user_is_rejected(cancellable_manager, coordination_backend): - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - - with ( - mock.patch("superset.extensions.celery_app") as celery_app, - raises(AsyncQueryTokenException), - ): - cancellable_manager.cancel_job("job-1", "chan-1", 999) - - celery_app.control.revoke.assert_not_called() - - -def test_cancel_job_wrong_channel_is_rejected( - cancellable_manager, coordination_backend -): - """A matching user on a different channel still cannot cancel the job.""" - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - - with ( - mock.patch("superset.extensions.celery_app") as celery_app, - raises(AsyncQueryTokenException), - ): - cancellable_manager.cancel_job("job-1", "other-chan", 7) - - celery_app.control.revoke.assert_not_called() - - -def test_cancel_job_unknown_raises(cancellable_manager, coordination_backend): - coordination_backend.get.return_value = None - - with ( - mock.patch("superset.extensions.celery_app") as celery_app, - raises(AsyncQueryJobException), - ): - cancellable_manager.cancel_job("job-1", "chan-1", 7) - - celery_app.control.revoke.assert_not_called() - - -def test_is_job_cancelled(cancellable_manager, coordination_backend): - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7, "cancelled": True} - ) - assert cancellable_manager.is_job_cancelled("job-1") is True - - coordination_backend.get.return_value = json.dumps( - {"channel_id": "chan-1", "user_id": 7} - ) - assert cancellable_manager.is_job_cancelled("job-1") is False - - coordination_backend.get.return_value = None - assert cancellable_manager.is_job_cancelled("job-1") is False - - -def test_is_job_cancelled_swallows_cache_errors( - cancellable_manager, coordination_backend -): - """A cache failure must not escape and mask the worker's original error.""" - coordination_backend.get.side_effect = RuntimeError("redis down") - assert cancellable_manager.is_job_cancelled("job-1") is False - - -def test_update_job_clears_registry_before_terminal_event( - cancellable_manager, coordination_backend -): - """Clearing first is what makes a cancel that lost the race a 404.""" - calls = [] - cancellable_manager._stream_limit = 100 - cancellable_manager._stream_limit_firehose = 1000 - coordination_backend.delete.side_effect = lambda *_: calls.append("delete") - coordination_backend.xadd.side_effect = lambda *_: calls.append("xadd") - job_metadata = {"channel_id": "chan-1", "job_id": "job-1", "user_id": 7} - - cancellable_manager.update_job(job_metadata, AsyncQueryManager.STATUS_DONE) - - coordination_backend.delete.assert_called_once_with("async-events-job-cancel:job-1") - assert calls == ["delete", "xadd", "xadd"] diff --git a/tests/unit_tests/charts/test_chart_data_api.py b/tests/unit_tests/charts/test_chart_data_api.py index 1c88bceb647..ab61735b363 100644 --- a/tests/unit_tests/charts/test_chart_data_api.py +++ b/tests/unit_tests/charts/test_chart_data_api.py @@ -587,15 +587,6 @@ def test_run_async_does_not_project_timing_onto_a_job_response( ) -> None: command = MagicMock() command.execute.side_effect = ChartDataCacheLoadError("cache miss") - async_command = MagicMock() - async_command.run.return_value = { - "channel_id": "channel", - "job_id": "job", - "user_id": 1, - "status": "pending", - "errors": [], - "result_url": "/api/v1/chart/data/job", - } api = ChartDataRestApi() original = app.config.get("CHART_DATA_INCLUDE_TIMING") @@ -604,9 +595,9 @@ def test_run_async_does_not_project_timing_onto_a_job_response( with ( app.test_request_context("/api/v1/chart/data", method="POST"), patch( - "superset.charts.data.api.CreateAsyncChartDataJobCommand", - return_value=async_command, - ), + "superset.charts.data.api.submit_chart_data_query_tasks", + return_value={"task_ids": ["task-1", "task-2"]}, + ) as submit, patch("superset.charts.data.api.get_user_id", return_value=1), ): response = api._run_async({"force": False}, command) @@ -614,8 +605,10 @@ def test_run_async_does_not_project_timing_onto_a_job_response( app.config["CHART_DATA_INCLUDE_TIMING"] = original assert response.status_code == 202 - assert "timing" not in json.loads(response.get_data(as_text=True)) - async_command.validate.assert_called_once() + body = json.loads(response.get_data(as_text=True)) + # The 202 body is just the async job (the query tasks to poll); no timing. + assert body == {"task_ids": ["task-1", "task-2"]} + submit.assert_called_once_with(command.query_context, 1) def test_run_async_projects_opt_in_timing_for_a_cached_result( diff --git a/tests/unit_tests/commands/chart/create_async_job_command_test.py b/tests/unit_tests/commands/chart/create_async_job_command_test.py deleted file mode 100644 index f723fcb5342..00000000000 --- a/tests/unit_tests/commands/chart/create_async_job_command_test.py +++ /dev/null @@ -1,51 +0,0 @@ -# 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. -from unittest.mock import MagicMock, patch - -import pytest - -from superset.commands.chart.data.create_async_job_command import ( - CreateAsyncChartDataJobCommand, -) - - -def test_run_without_validate_raises_clear_error(): - """run() must raise a clear error when validate() was not called first.""" - command = CreateAsyncChartDataJobCommand() - - with pytest.raises(RuntimeError, match="called before validate"): - command.run(form_data={}, user_id=1) - - -def test_run_after_validate_submits_job(): - """run() submits the job using the channel id captured during validate().""" - command = CreateAsyncChartDataJobCommand() - - with patch( - "superset.commands.chart.data.create_async_job_command.async_query_manager", - new=MagicMock(), - ) as mock_manager: - mock_manager.parse_channel_id_from_request.return_value = "channel-123" - mock_manager.submit_chart_data_job.return_value = {"job_id": "abc"} - - command.validate(request=MagicMock()) - result = command.run(form_data={"k": "v"}, user_id=42) - - mock_manager.submit_chart_data_job.assert_called_once_with( - "channel-123", {"k": "v"}, 42 - ) - assert result == {"job_id": "abc"} diff --git a/tests/unit_tests/coordination/test_service.py b/tests/unit_tests/coordination/test_service.py index f522a98eada..cc9419a2b84 100644 --- a/tests/unit_tests/coordination/test_service.py +++ b/tests/unit_tests/coordination/test_service.py @@ -53,26 +53,6 @@ def test_get_backend_none_when_distributed_coordination_unset( assert CoordinationService.is_backend_defined() is False -def test_get_backend_ignores_legacy_gaq_config( - app_context: None, mocker: MockerFixture -) -> None: - # The coordinator resolves DISTRIBUTED_COORDINATION_CONFIG only; the deprecated - # GAQ backend must never leak into locks/GTF, even with GAQ enabled. - _patch_distributed_coordination(mocker, None) - mocker.patch("superset.is_feature_enabled", return_value=True) - mocker.patch.dict( - "flask.current_app.config", - {"GLOBAL_ASYNC_QUERIES_CACHE_BACKEND": {"CACHE_TYPE": "RedisCache"}}, - ) - get_cache_backend = mocker.patch( - "superset.async_events.async_query_manager.get_cache_backend", - ) - - assert CoordinationService.get_backend() is None - assert CoordinationService.is_backend_defined() is False - get_cache_backend.assert_not_called() - - def test_backend_only_ops_raise_when_backend_unavailable( app_context: None, mocker: MockerFixture ) -> None: diff --git a/tests/unit_tests/initialization/check_async_query_secret_test.py b/tests/unit_tests/initialization/check_async_query_secret_test.py deleted file mode 100644 index 0675c215c9a..00000000000 --- a/tests/unit_tests/initialization/check_async_query_secret_test.py +++ /dev/null @@ -1,140 +0,0 @@ -# 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. -"""Unit tests for the default async JWT secret startup check.""" - -from unittest.mock import MagicMock - -import pytest -from pytest_mock import MockerFixture - -from superset.constants import CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET -from superset.initialization import SupersetAppInitializer - - -def _make_initializer( - secret: str, - *, - debug: bool = False, - testing: bool = False, -) -> SupersetAppInitializer: - initializer = SupersetAppInitializer.__new__(SupersetAppInitializer) - app = MagicMock() - app.debug = debug - app.config = { - "GLOBAL_ASYNC_QUERIES_JWT_SECRET": secret, - "TESTING": testing, - } - initializer.superset_app = app - initializer.config = app.config - return initializer - - -def test_check_async_query_secret_rejects_default_in_production( - mocker: MockerFixture, -) -> None: - """A default async secret with GAQ enabled refuses to start in production.""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=True, - ) - mocker.patch("superset.initialization.is_test", return_value=False) - initializer = _make_initializer(CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET) - - with pytest.raises(SystemExit): - initializer.check_async_query_secret() - - -def test_check_async_query_secret_allows_overridden_secret( - mocker: MockerFixture, -) -> None: - """A non-default async secret does not block startup.""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=True, - ) - mocker.patch("superset.initialization.is_test", return_value=False) - initializer = _make_initializer("a-strong-random-secret-value-1234567890") - - # Should not raise. - initializer.check_async_query_secret() - - -def test_check_async_query_secret_skipped_when_gaq_disabled( - mocker: MockerFixture, -) -> None: - """The check is a no-op when GLOBAL_ASYNC_QUERIES is disabled.""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=False, - ) - mocker.patch("superset.initialization.is_test", return_value=False) - initializer = _make_initializer(CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET) - - # Should not raise even with the default secret. - initializer.check_async_query_secret() - - -def test_check_async_query_secret_warns_only_in_debug( - mocker: MockerFixture, -) -> None: - """In debug the default secret warns but does not exit (matches SECRET_KEY).""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=True, - ) - mocker.patch("superset.initialization.is_test", return_value=False) - initializer = _make_initializer( - CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET, debug=True - ) - - # Should not raise in debug mode. - initializer.check_async_query_secret() - - -def test_configure_async_queries_skips_init_with_default_secret( - mocker: MockerFixture, -) -> None: - """In warn-only modes, async init is skipped so the short default secret cannot - crash startup via AsyncQueryManager.init_app()'s length check.""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=True, - ) - factory = mocker.patch("superset.initialization.async_query_manager_factory") - initializer = _make_initializer( - CHANGE_ME_GLOBAL_ASYNC_QUERIES_JWT_SECRET, debug=True - ) - - initializer.configure_async_queries() - - factory.init_app.assert_not_called() - - -def test_configure_async_queries_inits_with_overridden_secret( - mocker: MockerFixture, -) -> None: - """A non-default secret proceeds to initialize the async query manager.""" - mocker.patch( - "superset.initialization.feature_flag_manager.is_feature_enabled", - return_value=True, - ) - factory = mocker.patch("superset.initialization.async_query_manager_factory") - initializer = _make_initializer("a-strong-random-secret-value-1234567890") - - initializer.configure_async_queries() - - factory.init_app.assert_called_once_with(initializer.superset_app) diff --git a/tests/unit_tests/initialization/check_encryption_engine_test.py b/tests/unit_tests/initialization/check_encryption_engine_test.py index ff3ce93d585..2bdd154ceb3 100644 --- a/tests/unit_tests/initialization/check_encryption_engine_test.py +++ b/tests/unit_tests/initialization/check_encryption_engine_test.py @@ -76,7 +76,7 @@ def test_silent_when_engine_value_is_unrecognized() -> None: def test_never_raises_system_exit() -> None: - """Unlike check_secret_key/check_guest_token_secret/check_async_query_secret, + """Unlike check_secret_key/check_guest_token_secret, this check must never refuse to start: the legacy engine is a supported configuration, not a known-bad placeholder, so blocking startup on it would turn an opt-in hardening step into a forced-migration outage. diff --git a/tests/unit_tests/initialization_test.py b/tests/unit_tests/initialization_test.py index c34d2f157c9..fea1d86b32f 100644 --- a/tests/unit_tests/initialization_test.py +++ b/tests/unit_tests/initialization_test.py @@ -157,7 +157,6 @@ class TestSupersetAppInitializer: patch.object(app_initializer, "configure_url_map_converters"), patch.object(app_initializer, "configure_data_sources"), patch.object(app_initializer, "configure_auth_provider"), - patch.object(app_initializer, "configure_async_queries"), patch.object(app_initializer, "configure_ssh_manager"), patch.object(app_initializer, "configure_stats_manager"), patch.object(app_initializer, "init_views"), diff --git a/tests/unit_tests/tasks/test_async_queries.py b/tests/unit_tests/tasks/test_async_queries.py index 81863004803..83af4098228 100644 --- a/tests/unit_tests/tasks/test_async_queries.py +++ b/tests/unit_tests/tasks/test_async_queries.py @@ -14,292 +14,106 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import Any +"""Unit tests for the GTF chart-data fan-out orchestrator.""" + from unittest import mock +from uuid import uuid4 -import pytest -from celery.exceptions import SoftTimeLimitExceeded -from flask_babel import lazy_gettext as _ - -from superset.commands.chart.exceptions import ChartDataQueryFailedError -from superset.errors import ErrorLevel, SupersetError, SupersetErrorType -from superset.exceptions import ( - OAuth2RedirectError, - SupersetErrorException, - SupersetErrorsException, -) -from superset.utils.error_sanitization import GENERIC_ERROR_MESSAGE +from pytest_mock import MockerFixture -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_with_error( - mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager -): - """Test that the task is gracefully marked failed in event of error""" - from superset.tasks.async_queries import load_chart_data_into_cache +def _fake_query_context(num_queries: int, contribution_idx: int | None = None): + """Build a MagicMock QueryContext with ``num_queries`` queries. - job_metadata = {"user_id": 1} - form_data = {} - err_message = "Something went wrong" - err = ChartDataQueryFailedError(_(err_message)) - - mock_user = mock.MagicMock() - mock_query_context_schema = mock.MagicMock() - - mock_security_manager.get_user_by_id.return_value = mock_user - mock_async_query_manager.STATUS_ERROR = "error" - mock_query_context_schema_cls.return_value = mock_query_context_schema - - mock_query_context_schema.load.side_effect = err - - with pytest.raises(ChartDataQueryFailedError): - load_chart_data_into_cache(job_metadata, form_data) - - expected_errors = [{"message": err_message}] - - mock_async_query_manager.update_job.assert_called_once_with( - job_metadata, "error", errors=expected_errors - ) - - -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_cancelled_emits_no_event( - mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager -): - """A revoke leaves the terminal event to the cancel request that sent it.""" - from superset.tasks.async_queries import load_chart_data_into_cache - - job_metadata = {"user_id": 1, "job_id": "job-1"} - form_data: dict[str, Any] = {} - - mock_security_manager.get_user_by_id.return_value = mock.MagicMock() - # Sync Mock: is_job_cancelled is a plain method, but patching the manager - # yields an AsyncMock whose calls would otherwise return truthy coroutines. - mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=True) - mock_query_context_schema_cls.return_value.load.side_effect = ( - SoftTimeLimitExceeded() - ) - - with pytest.raises(SoftTimeLimitExceeded): - load_chart_data_into_cache(job_metadata, form_data) - - mock_async_query_manager.is_job_cancelled.assert_called_once_with("job-1") - mock_async_query_manager.update_job.assert_not_called() - - -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_timeout_emits_error( - mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager -): - """A genuine timeout reports an error, or the client waits forever.""" - from superset.tasks.async_queries import load_chart_data_into_cache - - job_metadata = {"user_id": 1, "job_id": "job-1"} - form_data: dict[str, Any] = {} - - mock_security_manager.get_user_by_id.return_value = mock.MagicMock() - mock_async_query_manager.STATUS_ERROR = "error" - mock_async_query_manager.is_job_cancelled = mock.Mock(return_value=False) - mock_query_context_schema_cls.return_value.load.side_effect = ( - SoftTimeLimitExceeded() - ) - - with pytest.raises(SoftTimeLimitExceeded): - load_chart_data_into_cache(job_metadata, form_data) - - mock_async_query_manager.update_job.assert_called_once_with( - job_metadata, - "error", - errors=[{"message": "A timeout occurred while loading chart data"}], - ) - - -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_with_superset_error_exception( - mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager -): - """Test that SupersetErrorException extracts SIP-40 style errors""" - from superset.tasks.async_queries import load_chart_data_into_cache - - job_metadata = {"user_id": 1} - form_data = {} - - superset_error = SupersetError( - message="Access denied to datasource", - error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR, - level=ErrorLevel.ERROR, - extra={"datasource": "my_table"}, - ) - err = SupersetErrorException(superset_error) - - mock_user = mock.MagicMock() - mock_query_context_schema = mock.MagicMock() - - mock_security_manager.get_user_by_id.return_value = mock_user - mock_async_query_manager.STATUS_ERROR = "error" - mock_query_context_schema_cls.return_value = mock_query_context_schema - - mock_query_context_schema.load.side_effect = err - - with pytest.raises(SupersetErrorException): - load_chart_data_into_cache(job_metadata, form_data) - - # Verify the full SIP-40 error structure is preserved - call_args = mock_async_query_manager.update_job.call_args - assert call_args[0] == (job_metadata, "error") - errors = call_args[1]["errors"] - assert len(errors) == 1 - assert errors[0]["message"] == "Access denied to datasource" - assert errors[0]["error_type"] == SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR - assert errors[0]["level"] == ErrorLevel.ERROR - assert errors[0]["extra"]["datasource"] == "my_table" - - -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_with_superset_errors_exception( - mock_query_context_schema_cls, mock_async_query_manager, mock_security_manager -): - """Test that SupersetErrorsException extracts multiple SIP-40 style errors""" - from superset.tasks.async_queries import load_chart_data_into_cache - - job_metadata = {"user_id": 1} - form_data = {} - - superset_errors = [ - SupersetError( - message="Column not found", - error_type=SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR, - level=ErrorLevel.ERROR, - ), - SupersetError( - message="Table not found", - error_type=SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR, - level=ErrorLevel.WARNING, - ), - ] - err = SupersetErrorsException(superset_errors) - - mock_user = mock.MagicMock() - mock_query_context_schema = mock.MagicMock() - - mock_security_manager.get_user_by_id.return_value = mock_user - mock_async_query_manager.STATUS_ERROR = "error" - mock_query_context_schema_cls.return_value = mock_query_context_schema - - mock_query_context_schema.load.side_effect = err - - with pytest.raises(SupersetErrorsException): - load_chart_data_into_cache(job_metadata, form_data) - - # Verify all SIP-40 errors are preserved - call_args = mock_async_query_manager.update_job.call_args - assert call_args[0] == (job_metadata, "error") - errors = call_args[1]["errors"] - assert len(errors) == 2 - assert errors[0]["message"] == "Column not found" - assert errors[0]["error_type"] == SupersetErrorType.COLUMN_DOES_NOT_EXIST_ERROR - assert errors[0]["level"] == ErrorLevel.ERROR - assert errors[1]["message"] == "Table not found" - assert errors[1]["error_type"] == SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR - assert errors[1]["level"] == ErrorLevel.WARNING - - -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.commands.chart.data.get_data_command.ChartDataCommand") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_preserves_oauth2_redirect_error( - mock_query_context_schema_cls: mock.MagicMock, - mock_command_cls: mock.MagicMock, - mock_async_query_manager: mock.MagicMock, - mock_security_manager: mock.MagicMock, -) -> None: + When ``contribution_idx`` is set, ``prepare_contribution_totals`` reports that + query as contribution-coupled with query 0 as the totals query. """ - OAuth2RedirectError raised by ``ChartDataCommand.run`` must reach the async - job's errors list as a structured SIP-40 envelope (with ``error_type`` and - the OAuth2 ``extra`` payload) instead of being flattened to a plain - message, so dashboard charts can render the OAuth2 banner when - GLOBAL_ASYNC_QUERIES is enabled. - """ - from superset.tasks.async_queries import load_chart_data_into_cache + ctx = mock.MagicMock() + ctx.queries = [mock.MagicMock(name=f"q{i}") for i in range(num_queries)] + ctx.cache_values = {"queries": [{"i": i} for i in range(num_queries)]} + ctx.query_cache_key.side_effect = lambda q: f"key-{ctx.queries.index(q)}" + if contribution_idx is not None: + ctx.prepare_contribution_totals.return_value = ([contribution_idx], 0) + else: + ctx.prepare_contribution_totals.return_value = ([], None) + return ctx - job_metadata = {"user_id": 1} - form_data: dict[str, Any] = {} - mock_security_manager.get_user_by_id.return_value = mock.MagicMock() - mock_async_query_manager.STATUS_ERROR = "error" - mock_query_context_schema_cls.return_value.load.return_value = mock.MagicMock() +def _patch_schedule(mocker: MockerFixture): + """Patch execute_chart_query.schedule to return Tasks with unique uuids.""" + scheduled = [] - mock_command_cls.return_value.run.side_effect = OAuth2RedirectError( - url="https://accounts.example.com/o/oauth2/v2/auth?...", - tab_id="tab-123", - redirect_uri="https://superset.example.com/oauth2/redirect", + def _schedule(*args, **kwargs): + task = mock.MagicMock() + task.uuid = uuid4() + scheduled.append({"args": args, "kwargs": kwargs, "task": task}) + return task + + mocker.patch( + "superset.tasks.async_queries.execute_chart_query.schedule", + side_effect=_schedule, ) - - with pytest.raises(OAuth2RedirectError): - load_chart_data_into_cache(job_metadata, form_data) - - call_args = mock_async_query_manager.update_job.call_args - assert call_args[0] == (job_metadata, "error") - errors = call_args[1]["errors"] - assert len(errors) == 1 - # A flattened error would only carry the generic permission message; the - # structured envelope must be preserved for the frontend OAuth2 banner. - assert errors[0] != {"message": "You don't have permission to access the data."} - assert errors[0]["error_type"] == SupersetErrorType.OAUTH2_REDIRECT - assert errors[0]["level"] == ErrorLevel.WARNING - assert errors[0]["extra"] == { - "url": "https://accounts.example.com/o/oauth2/v2/auth?...", - "tab_id": "tab-123", - "redirect_uri": "https://superset.example.com/oauth2/redirect", - } - - -@mock.patch("superset.security.SupersetSecurityManager.is_guest_user") -@mock.patch("superset.tasks.async_queries.security_manager") -@mock.patch("superset.tasks.async_queries.async_query_manager") -@mock.patch("superset.tasks.async_queries.ChartDataQueryContextSchema") -def test_load_chart_data_into_cache_redacts_error_for_guest_user( - mock_query_context_schema_cls, - mock_async_query_manager, - mock_security_manager, - mock_is_guest_user, - app_context: None, -): - """An embedded viewer gets a generic message instead of the engine's.""" - from superset.tasks.async_queries import load_chart_data_into_cache - - job_metadata = {"user_id": 1} - err = SupersetErrorsException( - [ - SupersetError( - message="Table mydb.myschema.mytable was not found", - error_type=SupersetErrorType.TABLE_DOES_NOT_EXIST_ERROR, - level=ErrorLevel.ERROR, - extra={"engine_name": "BigQuery"}, - ) - ] + mocker.patch( + "superset.tasks.async_queries.serialize_query", + side_effect=lambda ctx, index: {"query": index}, ) + guest = mocker.patch("superset.tasks.async_queries.security_manager") + # Force a sync return (a bare patched method resolves to an AsyncMock whose + # call is a truthy coroutine here); see the GAQ→GTF testing notes. + guest.get_current_guest_user_if_guest = mock.MagicMock(return_value=None) + return scheduled - mock_is_guest_user.return_value = True - mock_security_manager.get_user_by_id.return_value = mock.MagicMock() - mock_async_query_manager.STATUS_ERROR = "error" - mock_query_context_schema_cls.return_value.load.side_effect = err - with pytest.raises(SupersetErrorsException): - load_chart_data_into_cache(job_metadata, {}) +def test_fan_out_schedules_one_task_per_query(mocker: MockerFixture) -> None: + from superset.tasks.async_queries import submit_chart_data_query_tasks - errors = mock_async_query_manager.update_job.call_args[1]["errors"] - assert errors[0]["message"] == str(GENERIC_ERROR_MESSAGE) - assert errors[0]["error_type"] == SupersetErrorType.GENERIC_BACKEND_ERROR - assert "engine_name" not in errors[0]["extra"] + scheduled = _patch_schedule(mocker) + ctx = _fake_query_context(3) + + result = submit_chart_data_query_tasks(ctx, user_id=7) + + assert len(scheduled) == 3 + # The 202 body carries the query tasks' uuids, in query order. + assert result["task_ids"] == [str(s["task"].uuid) for s in scheduled] + # Independent queries carry no dependency and no totals key. + for call in scheduled: + assert call["kwargs"]["options"].depends_on is None + assert call["args"][3] is None # totals_cache_key + + +def test_contribution_query_depends_on_totals(mocker: MockerFixture) -> None: + from superset.tasks.async_queries import submit_chart_data_query_tasks + + scheduled = _patch_schedule(mocker) + # query 1 is a contribution query; query 0 is the totals query. + ctx = _fake_query_context(2, contribution_idx=1) + + submit_chart_data_query_tasks(ctx, user_id=7) + + # Totals query (index 0) is scheduled first with no dependency. + totals_call = scheduled[0] + assert totals_call["kwargs"]["options"].depends_on is None + # The totals query's row_limit is normalized so its key matches the entry + # its dependents read. + assert ctx.cache_values["queries"][0]["row_limit"] is None + + # The contribution query depends on the totals task and receives its key. + dep_call = next(c for c in scheduled if c["args"][0] == {"query": 1}) + assert dep_call["kwargs"]["options"].depends_on == [totals_call["task"]] + assert dep_call["args"][3] == "key-0" # totals_cache_key + + +def test_guest_token_forwarded(mocker: MockerFixture) -> None: + from superset.tasks.async_queries import submit_chart_data_query_tasks + + scheduled = _patch_schedule(mocker) + guest = mock.MagicMock() + guest.guest_token = {"user": {"username": "guest"}} + sm = mocker.patch("superset.tasks.async_queries.security_manager") + sm.get_current_guest_user_if_guest = mock.MagicMock(return_value=guest) + ctx = _fake_query_context(1) + + submit_chart_data_query_tasks(ctx, user_id=None) + + # The guest token is passed to the task so the worker can impersonate. + assert scheduled[0]["args"][2] == guest.guest_token