From f4bdcb5743d7f70048d922500975496f8f219dc7 Mon Sep 17 00:00:00 2001 From: "JUST.in DO IT" Date: Mon, 18 Mar 2024 12:52:23 -0700 Subject: [PATCH] perf(sqllab): reduce bootstrap data delay by queries (#27488) --- .../QueryHistory/QueryHistory.test.tsx | 73 +++++++- .../SqlLab/components/QueryHistory/index.tsx | 106 ++++++++--- .../SqlLab/reducers/getInitialState.test.ts | 21 ++- .../src/SqlLab/reducers/getInitialState.ts | 7 +- .../src/hooks/apiResources/queries.test.ts | 154 +++++++++++++++ .../src/hooks/apiResources/queries.ts | 176 ++++++++++++++++++ .../src/hooks/apiResources/queryApi.ts | 3 +- superset/queries/api.py | 19 +- superset/sqllab/utils.py | 15 +- tests/integration_tests/sql_lab/api_tests.py | 22 +-- tests/integration_tests/sqllab_tests.py | 57 ++++++ 11 files changed, 577 insertions(+), 76 deletions(-) create mode 100644 superset-frontend/src/hooks/apiResources/queries.test.ts create mode 100644 superset-frontend/src/hooks/apiResources/queries.ts diff --git a/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx b/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx index ad1881b5d93ed..110e7b4ae2f20 100644 --- a/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx +++ b/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx @@ -17,7 +17,10 @@ * under the License. */ import React from 'react'; -import { render, screen } from 'spec/helpers/testing-library'; +import fetchMock from 'fetch-mock'; +import * as uiCore from '@superset-ui/core'; +import { FeatureFlag, QueryState } from '@superset-ui/core'; +import { render, screen, waitFor } from 'spec/helpers/testing-library'; import QueryHistory from 'src/SqlLab/components/QueryHistory'; import { initialState } from 'src/SqlLab/fixtures'; @@ -27,18 +30,72 @@ const mockedProps = { latestQueryId: 'yhMUZCGb', }; +const fakeApiResult = { + count: 4, + ids: [692], + result: [ + { + changed_on: '2024-03-12T20:01:02.497775', + client_id: 'b0ZDzRYzn', + database: { + database_name: 'examples', + id: 1, + }, + end_time: '1710273662496.047852', + error_message: null, + executed_sql: 'SELECT * from "FCC 2018 Survey"\nLIMIT 1001', + id: 692, + limit: 1000, + limiting_factor: 'DROPDOWN', + progress: 100, + results_key: null, + rows: 443, + schema: 'main', + select_as_cta: false, + sql: 'SELECT * from "FCC 2018 Survey" ', + sql_editor_id: '22', + start_time: '1710273662445.992920', + status: QueryState.Success, + tab_name: 'Untitled Query 16', + tmp_table_name: null, + tracking_url: null, + user: { + first_name: 'admin', + id: 1, + last_name: 'user', + }, + }, + ], +}; + const setup = (overrides = {}) => ( ); -describe('QueryHistory', () => { - it('Renders an empty state for query history', () => { - render(setup(), { useRedux: true, initialState }); +test('Renders an empty state for query history', () => { + render(setup(), { useRedux: true, initialState }); + + const emptyStateText = screen.getByText( + /run a query to display query history/i, + ); + + expect(emptyStateText).toBeVisible(); +}); - const emptyStateText = screen.getByText( - /run a query to display query history/i, +test('fetches the query history when the persistence mode is enabled', async () => { + const isFeatureEnabledMock = jest + .spyOn(uiCore, 'isFeatureEnabled') + .mockImplementation( + featureFlag => featureFlag === FeatureFlag.SqllabBackendPersistence, ); - expect(emptyStateText).toBeVisible(); - }); + const editorQueryApiRoute = `glob:*/api/v1/query/?q=*`; + fetchMock.get(editorQueryApiRoute, fakeApiResult); + render(setup(), { useRedux: true, initialState }); + await waitFor(() => + expect(fetchMock.calls(editorQueryApiRoute).length).toBe(1), + ); + const queryResultText = screen.getByText(fakeApiResult.result[0].rows); + expect(queryResultText).toBeInTheDocument(); + isFeatureEnabledMock.mockClear(); }); diff --git a/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx b/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx index 311a125d556c3..e020c3302ff38 100644 --- a/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx +++ b/superset-frontend/src/SqlLab/components/QueryHistory/index.tsx @@ -16,12 +16,23 @@ * specific language governing permissions and limitations * under the License. */ -import React, { useMemo } from 'react'; +import React, { useEffect, useMemo, useState } from 'react'; import { shallowEqual, useSelector } from 'react-redux'; +import { useInView } from 'react-intersection-observer'; +import { omit } from 'lodash'; import { EmptyStateMedium } from 'src/components/EmptyState'; -import { t, styled } from '@superset-ui/core'; +import { + t, + styled, + css, + FeatureFlag, + isFeatureEnabled, +} from '@superset-ui/core'; import QueryTable from 'src/SqlLab/components/QueryTable'; import { SqlLabRootState } from 'src/SqlLab/types'; +import { useEditorQueriesQuery } from 'src/hooks/apiResources/queries'; +import { Skeleton } from 'src/components'; +import useEffectEvent from 'src/hooks/useEffectEvent'; interface QueryHistoryProps { queryEditorId: string | number; @@ -40,39 +51,92 @@ const StyledEmptyStateWrapper = styled.div` } `; +const getEditorQueries = ( + queries: SqlLabRootState['sqlLab']['queries'], + queryEditorId: string | number, +) => + Object.values(queries).filter( + ({ sqlEditorId }) => String(sqlEditorId) === String(queryEditorId), + ); + const QueryHistory = ({ queryEditorId, displayLimit, latestQueryId, }: QueryHistoryProps) => { + const [ref, hasReachedBottom] = useInView({ threshold: 0 }); + const [pageIndex, setPageIndex] = useState(0); const queries = useSelector( ({ sqlLab: { queries } }: SqlLabRootState) => queries, shallowEqual, ); + const { data, isLoading, isFetching } = useEditorQueriesQuery( + { editorId: `${queryEditorId}`, pageIndex }, + { + skip: !isFeatureEnabled(FeatureFlag.SqllabBackendPersistence), + }, + ); const editorQueries = useMemo( () => - Object.values(queries).filter( - ({ sqlEditorId }) => String(sqlEditorId) === String(queryEditorId), - ), - [queries, queryEditorId], + data + ? getEditorQueries( + omit( + queries, + data.result.map(({ id }) => id), + ), + queryEditorId, + ) + .concat(data.result) + .reverse() + : getEditorQueries(queries, queryEditorId), + [queries, data, queryEditorId], ); + const loadNext = useEffectEvent(() => { + setPageIndex(pageIndex + 1); + }); + + const loadedDataCount = data?.result.length || 0; + const totalCount = data?.count || 0; + + useEffect(() => { + if (hasReachedBottom && loadedDataCount < totalCount) { + loadNext(); + } + }, [hasReachedBottom, loadNext, loadedDataCount, totalCount]); + + if (!editorQueries.length && isLoading) { + return ; + } + return editorQueries.length > 0 ? ( - + <> + + {data && loadedDataCount < totalCount && ( +
+ )} + {isFetching && } + ) : ( { }), ); + const latestQuery = { + ...runningQuery, + id: 'latestPersisted', + startDttm: Number(startDttmInStr), + endDttm: Number(endDttmInStr), + }; const initializedQueries = getInitialState({ - ...apiData, - queries: { - backendPersisted: { - ...runningQuery, - id: 'backendPersisted', - startDttm: startDttmInStr, - endDttm: endDttmInStr, - }, + ...apiDataWithTabState, + active_tab: { + ...apiDataWithTabState.active_tab, + latest_query: latestQuery, }, }).sqlLab.queries; - expect(initializedQueries.backendPersisted).toEqual( + expect(initializedQueries.latestPersisted).toEqual( expect.objectContaining({ startDttm: Number(startDttmInStr), endDttm: Number(endDttmInStr), diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.ts index 04e7f4eca4573..bcdc1f40c33da 100644 --- a/superset-frontend/src/SqlLab/reducers/getInitialState.ts +++ b/superset-frontend/src/SqlLab/reducers/getInitialState.ts @@ -136,7 +136,12 @@ export default function getInitialState({ }); } - const queries = { ...queries_ }; + const queries = { + ...queries_, + ...(activeTab?.latest_query && { + [activeTab.latest_query.id]: activeTab.latest_query, + }), + }; /** * If the `SQLLAB_BACKEND_PERSISTENCE` feature flag is off, or if the user diff --git a/superset-frontend/src/hooks/apiResources/queries.test.ts b/superset-frontend/src/hooks/apiResources/queries.test.ts new file mode 100644 index 0000000000000..14d1f9ceacd22 --- /dev/null +++ b/superset-frontend/src/hooks/apiResources/queries.test.ts @@ -0,0 +1,154 @@ +/** + * 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 rison from 'rison'; +import fetchMock from 'fetch-mock'; +import { act, renderHook } from '@testing-library/react-hooks'; +import { + createWrapper, + defaultStore as store, +} from 'spec/helpers/testing-library'; +import { QueryState } from '@superset-ui/core'; +import { api } from 'src/hooks/apiResources/queryApi'; +import { mapQueryResponse, useEditorQueriesQuery } from './queries'; + +const fakeApiResult = { + count: 4, + ids: [692], + result: [ + { + changed_on: '2024-03-12T20:01:02.497775', + client_id: 'b0ZDzRYzn', + database: { + database_name: 'examples', + id: 1, + }, + end_time: '1710273662496.047852', + error_message: null, + executed_sql: 'SELECT * from "FCC 2018 Survey"\nLIMIT 1001', + id: 692, + limit: 1000, + limiting_factor: 'DROPDOWN', + progress: 100, + results_key: null, + rows: 1000, + schema: 'main', + select_as_cta: false, + sql: 'SELECT * from "FCC 2018 Survey" ', + sql_editor_id: '22', + start_time: '1710273662445.992920', + status: QueryState.Success, + tab_name: 'Untitled Query 16', + tmp_table_name: null, + tracking_url: null, + user: { + first_name: 'admin', + id: 1, + last_name: 'user', + }, + }, + ], +}; + +afterEach(() => { + fetchMock.reset(); + act(() => { + store.dispatch(api.util.resetApiState()); + }); +}); + +test('returns api response mapping camelCase keys', async () => { + const editorId = '23'; + const editorQueryApiRoute = `glob:*/api/v1/query/?q=*`; + fetchMock.get(editorQueryApiRoute, fakeApiResult); + const { result, waitFor } = renderHook( + () => useEditorQueriesQuery({ editorId }), + { + wrapper: createWrapper({ + useRedux: true, + store, + }), + }, + ); + await waitFor(() => + expect(fetchMock.calls(editorQueryApiRoute).length).toBe(1), + ); + const expectedResult = { + ...fakeApiResult, + result: fakeApiResult.result.map(mapQueryResponse), + }; + expect( + rison.decode(fetchMock.calls(editorQueryApiRoute)[0][0].split('?q=')[1]), + ).toEqual( + expect.objectContaining({ + order_column: 'start_time', + order_direction: 'desc', + page: 0, + page_size: 25, + filters: [ + { + col: 'sql_editor_id', + opr: 'eq', + value: expect.stringContaining(editorId), + }, + ], + }), + ); + expect(result.current.data).toEqual(expectedResult); +}); + +test('merges paginated results', async () => { + const editorId = '23'; + const editorQueryApiRoute = `glob:*/api/v1/query/?q=*`; + fetchMock.get(editorQueryApiRoute, fakeApiResult); + const { waitFor } = renderHook(() => useEditorQueriesQuery({ editorId }), { + wrapper: createWrapper({ + useRedux: true, + store, + }), + }); + await waitFor(() => + expect(fetchMock.calls(editorQueryApiRoute).length).toBe(1), + ); + const { result: paginatedResult } = renderHook( + () => useEditorQueriesQuery({ editorId, pageIndex: 1 }), + { + wrapper: createWrapper({ + useRedux: true, + store, + }), + }, + ); + await waitFor(() => + expect(fetchMock.calls(editorQueryApiRoute).length).toBe(2), + ); + expect( + rison.decode(fetchMock.calls(editorQueryApiRoute)[1][0].split('?q=')[1]), + ).toEqual( + expect.objectContaining({ + page: 1, + }), + ); + expect(paginatedResult.current.data).toEqual({ + ...fakeApiResult, + result: [ + ...fakeApiResult.result.map(mapQueryResponse), + ...fakeApiResult.result.map(mapQueryResponse), + ], + }); +}); diff --git a/superset-frontend/src/hooks/apiResources/queries.ts b/superset-frontend/src/hooks/apiResources/queries.ts new file mode 100644 index 0000000000000..ad7ed8eb591c6 --- /dev/null +++ b/superset-frontend/src/hooks/apiResources/queries.ts @@ -0,0 +1,176 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import type { Query, QueryResponse } from '@superset-ui/core'; +import type { JsonResponse } from './queryApi'; +import { api } from './queryApi'; + +export type QueryResult = { + count: number; + ids: Query['id'][]; + result: QueryResponse[]; +}; + +export type EditorQueriesParams = { + editorId: string; + pageIndex?: number; + pageSize?: number; +}; + +interface ResponseResult { + id: Query['queryId']; + client_id: Query['id']; + database: { + id: Query['dbId']; + database_name: string; + }; + executed_sql: Query['executedSql']; + error_message: Query['errorMessage']; + limit: Query['queryLimit']; + limiting_factor: Query['limitingFactor']; + progress: Query['progress']; + rows: Query['rows']; + select_as_cta: Query['ctas']; + schema: Query['schema']; + sql: Query['sql']; + sql_editor_id: Query['sqlEditorId']; + status: Query['state']; + tab_name: Query['tab']; + user: { + id: Query['userId']; + }; + start_time: string; + end_time: string; + tmp_table_name: Query['tempTable'] | null; + tracking_url: Query['trackingUrl']; + results_key: Query['resultsKey']; +} + +export const mapQueryResponse = ( + query: ResponseResult, +): Omit< + Query, + | 'tempSchema' + | 'started' + | 'time' + | 'duration' + | 'templateParams' + | 'querylink' + | 'output' + | 'actions' + | 'type' + | 'columns' +> => ({ + queryId: query.id, + id: query.client_id, + dbId: query.database.id, + db: query.database, + executedSql: query.executed_sql, + errorMessage: query.error_message, + queryLimit: query.limit, + ctas: query.select_as_cta, + limitingFactor: query.limiting_factor, + progress: query.progress, + rows: query.rows, + schema: query.schema, + sql: query.sql, + sqlEditorId: query.sql_editor_id, + state: query.status, + tab: query.tab_name, + startDttm: Number(query.start_time), + endDttm: Number(query.end_time), + tempTable: query.tmp_table_name || '', + trackingUrl: query.tracking_url, + resultsKey: query.results_key, + userId: query.user.id, + cached: false, + extra: { + progress: null, + }, + isDataPreview: false, + user: query.user, +}); + +const queryHistoryApi = api.injectEndpoints({ + endpoints: builder => ({ + editorQueries: builder.query({ + providesTags: ['EditorQueries'], + query: ({ editorId, pageIndex = 0, pageSize = 25 }) => ({ + method: 'GET', + endpoint: `/api/v1/query/`, + urlParams: { + keys: ['none'], + columns: [ + 'id', + 'client_id', + 'changed_on', + 'database.id', + 'database.database_name', + 'executed_sql', + 'error_message', + 'limit', + 'limiting_factor', + 'progress', + 'rows', + 'select_as_cta', + 'schema', + 'sql', + 'sql_editor_id', + 'status', + 'tab_name', + 'user.first_name', + 'user.id', + 'user.last_name', + 'start_time', + 'end_time', + 'tmp_table_name', + 'tmp_schema_name', + 'tracking_url', + 'results_key', + ], + order_column: 'start_time', + order_direction: 'desc', + page: pageIndex, + page_size: pageSize, + filters: [ + { + col: 'sql_editor_id', + opr: 'eq', + value: editorId, + }, + ], + }, + headers: { 'Content-Type': 'application/json' }, + transformResponse: ({ json }: JsonResponse) => ({ + ...json, + result: json.result.map(mapQueryResponse), + }), + }), + serializeQueryArgs: ({ queryArgs: { editorId } }) => ({ editorId }), + // Refetch when the page arg changes + forceRefetch({ currentArg, previousArg }) { + return currentArg !== previousArg; + }, + merge: (currentCache, newItems) => { + currentCache.result.push(...newItems.result); + }, + }), + }), +}); + +export const { useEditorQueriesQuery } = queryHistoryApi; diff --git a/superset-frontend/src/hooks/apiResources/queryApi.ts b/superset-frontend/src/hooks/apiResources/queryApi.ts index 4099422b5468f..ae25f3cc56d1a 100644 --- a/superset-frontend/src/hooks/apiResources/queryApi.ts +++ b/superset-frontend/src/hooks/apiResources/queryApi.ts @@ -37,7 +37,7 @@ export const supersetClientQuery: BaseQueryFn< endpoint: string; parseMethod?: ParseMethod; transformResponse?: (response: SupersetClientResponse) => JsonValue; - urlParams?: Record; + urlParams?: Record; }, JsonValue, ClientErrorObject @@ -80,6 +80,7 @@ export const api = createApi({ 'QueryValidations', 'TableMetadatas', 'SqlLabInitialState', + 'EditorQueries', ], endpoints: () => ({}), baseQuery: supersetClientQuery, diff --git a/superset/queries/api.py b/superset/queries/api.py index 9568542f76f0b..0695946fe07f3 100644 --- a/superset/queries/api.py +++ b/superset/queries/api.py @@ -71,11 +71,19 @@ class QueryRestApi(BaseSupersetModelRestApi): list_columns = [ "id", "changed_on", + "client_id", + "database.id", "database.database_name", "executed_sql", + "error_message", + "limit", + "limiting_factor", + "progress", "rows", "schema", + "select_as_cta", "sql", + "sql_editor_id", "sql_tables", "status", "tab_name", @@ -86,6 +94,7 @@ class QueryRestApi(BaseSupersetModelRestApi): "end_time", "tmp_table_name", "tracking_url", + "results_key", ] show_columns = [ "id", @@ -143,7 +152,15 @@ class QueryRestApi(BaseSupersetModelRestApi): "user": RelatedFieldFilter("first_name", FilterRelatedOwners), } - search_columns = ["changed_on", "database", "sql", "status", "user", "start_time"] + search_columns = [ + "changed_on", + "database", + "sql", + "status", + "user", + "start_time", + "sql_editor_id", + ] allowed_rel_fields = {"database", "user"} allowed_distinct_fields = {"status"} diff --git a/superset/sqllab/utils.py b/superset/sqllab/utils.py index 989bd19cc31c9..ca331d1e3459f 100644 --- a/superset/sqllab/utils.py +++ b/superset/sqllab/utils.py @@ -23,7 +23,7 @@ from superset import db, is_feature_enabled from superset.common.db_query_status import QueryStatus from superset.daos.database import DatabaseDAO -from superset.models.sql_lab import Query, TabState +from superset.models.sql_lab import TabState DATABASE_KEYS = [ "allow_file_upload", @@ -87,7 +87,6 @@ def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]: k: v for k, v in database.to_json().items() if k in DATABASE_KEYS } databases[database.id]["backend"] = database.backend - queries: dict[str, Any] = {} # These are unnecessary if sqllab backend persistence is disabled if is_feature_enabled("SQLLAB_BACKEND_PERSISTENCE"): @@ -97,7 +96,6 @@ def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]: .filter_by(user_id=user_id) .all() ) - tab_state_ids = [str(tab_state[0]) for tab_state in tabs_state] # return first active tab, or fallback to another one if no tab is active active_tab = ( db.session.query(TabState) @@ -105,20 +103,9 @@ def bootstrap_sqllab_data(user_id: int | None) -> dict[str, Any]: .order_by(TabState.active.desc()) .first() ) - # return all user queries associated with existing SQL editors - user_queries = ( - db.session.query(Query) - .filter_by(user_id=user_id) - .filter(Query.sql_editor_id.in_(tab_state_ids)) - .all() - ) - queries = { - query.client_id: dict(query.to_dict().items()) for query in user_queries - } return { "tab_state_ids": tabs_state, "active_tab": active_tab.to_dict() if active_tab else None, "databases": databases, - "queries": queries, } diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index d47bff92f72cd..5621ec0c3dffc 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -57,7 +57,6 @@ def test_get_from_empty_bootsrap_data(self): data = json.loads(resp.data.decode("utf-8")) result = data.get("result") assert result["active_tab"] == None - assert result["queries"] == {} assert result["tab_state_ids"] == [] self.assertEqual(len(result["databases"]), 0) @@ -87,7 +86,6 @@ def test_get_from_bootstrap_data_for_non_persisted_tab_state(self): data = json.loads(resp.data.decode("utf-8")) result = data.get("result") assert result["active_tab"] == None - assert result["queries"] == {} assert result["tab_state_ids"] == [] @mock.patch.dict( @@ -95,7 +93,7 @@ def test_get_from_bootstrap_data_for_non_persisted_tab_state(self): {"SQLLAB_BACKEND_PERSISTENCE": True}, clear=True, ) - def test_get_from_bootstrap_data_with_queries(self): + def test_get_from_bootstrap_data_with_latest_query(self): username = "admin" self.login(username) @@ -115,27 +113,11 @@ def test_get_from_bootstrap_data_with_queries(self): resp = self.get_json_resp("/tabstateview/", data=data) tab_state_id = resp["id"] - # run a query in the created tab - self.run_sql( - "SELECT name FROM birth_names", - "client_id_1", - username=username, - raise_on_error=True, - sql_editor_id=str(tab_state_id), - ) - # run an orphan query (no tab) - self.run_sql( - "SELECT name FROM birth_names", - "client_id_2", - username=username, - raise_on_error=True, - ) - # we should have only 1 query returned, since the second one is not # associated with any tabs resp = self.get_json_resp("/api/v1/sqllab/") result = resp["result"] - self.assertEqual(len(result["queries"]), 1) + self.assertEqual(result["active_tab"]["id"], tab_state_id) @mock.patch.dict( "superset.extensions.feature_flag_manager._feature_flags", diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py index 0dc4e26acad91..30b8401cc6f6d 100644 --- a/tests/integration_tests/sqllab_tests.py +++ b/tests/integration_tests/sqllab_tests.py @@ -461,6 +461,63 @@ def test_query_api_can_access_all_queries(self) -> None: db.session.commit() + def test_query_api_can_access_sql_editor_id_associated_queries(self) -> None: + """ + Test query api with sql_editor_id filter to + gamma and make sure sql editor associated queries show up. + """ + username = "gamma_sqllab" + self.login("gamma_sqllab") + + # create a tab + data = { + "queryEditor": json.dumps( + { + "title": "Untitled Query 1", + "dbId": 1, + "schema": None, + "autorun": False, + "sql": "SELECT ...", + "queryLimit": 1000, + } + ) + } + resp = self.get_json_resp("/tabstateview/", data=data) + tab_state_id = resp["id"] + # run a query in the created tab + self.run_sql( + "SELECT 1", + "client_id_1", + username=username, + raise_on_error=True, + sql_editor_id=str(tab_state_id), + ) + self.run_sql( + "SELECT 2", + "client_id_2", + username=username, + raise_on_error=True, + sql_editor_id=str(tab_state_id), + ) + # run an orphan query (no tab) + self.run_sql( + "SELECT 3", + "client_id_3", + username=username, + raise_on_error=True, + ) + + arguments = { + "filters": [ + {"col": "sql_editor_id", "opr": "eq", "value": str(tab_state_id)} + ] + } + url = f"/api/v1/query/?q={prison.dumps(arguments)}" + self.assertEqual( + {"SELECT 1", "SELECT 2"}, + {r.get("sql") for r in self.get_json_resp(url)["result"]}, + ) + def test_query_admin_can_access_all_queries(self) -> None: """ Test query api with all_query_access perm added to