You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2024/03/19 13:09:59 UTC
(superset) 06/06: perf(sqllab): reduce bootstrap data delay by queries (#27488)
This is an automated email from the ASF dual-hosted git repository.
michaelsmolina pushed a commit to branch 4.0
in repository https://gitbox.apache.org/repos/asf/superset.git
commit f52647ff48219482bc65191834e9815a7e1f6dc3
Author: JUST.in DO IT <ju...@airbnb.com>
AuthorDate: Mon Mar 18 12:52:23 2024 -0700
perf(sqllab): reduce bootstrap data delay by queries (#27488)
(cherry picked from commit f4bdcb5743d7f70048d922500975496f8f219dc7)
---
.../components/QueryHistory/QueryHistory.test.tsx | 73 ++++++++-
.../src/SqlLab/components/QueryHistory/index.tsx | 106 ++++++++++---
.../src/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(-)
diff --git a/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx b/superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx
index ad1881b5d9..110e7b4ae2 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 = {}) => (
<QueryHistory {...mockedProps} {...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 311a125d55..e020c3302f 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 <Skeleton active />;
+ }
+
return editorQueries.length > 0 ? (
- <QueryTable
- columns={[
- 'state',
- 'started',
- 'duration',
- 'progress',
- 'rows',
- 'sql',
- 'results',
- 'actions',
- ]}
- queries={editorQueries}
- displayLimit={displayLimit}
- latestQueryId={latestQueryId}
- />
+ <>
+ <QueryTable
+ columns={[
+ 'state',
+ 'started',
+ 'duration',
+ 'progress',
+ 'rows',
+ 'sql',
+ 'results',
+ 'actions',
+ ]}
+ queries={editorQueries}
+ displayLimit={displayLimit}
+ latestQueryId={latestQueryId}
+ />
+ {data && loadedDataCount < totalCount && (
+ <div
+ ref={ref}
+ css={css`
+ position: relative;
+ top: -150px;
+ `}
+ />
+ )}
+ {isFetching && <Skeleton active />}
+ </>
) : (
<StyledEmptyStateWrapper>
<EmptyStateMedium
diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts
index 1dd3220fcc..6d6e65bad3 100644
--- a/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts
+++ b/superset-frontend/src/SqlLab/reducers/getInitialState.test.ts
@@ -25,7 +25,6 @@ const apiData = {
common: DEFAULT_COMMON_BOOTSTRAP_DATA,
tab_state_ids: [],
databases: [],
- queries: {},
user: {
userId: 1,
username: 'some name',
@@ -220,18 +219,20 @@ describe('getInitialState', () => {
}),
);
+ 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 04e7f4eca4..bcdc1f40c3 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 0000000000..14d1f9ceac
--- /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 0000000000..ad7ed8eb59
--- /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<QueryResult, EditorQueriesParams>({
+ 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 4099422b54..ae25f3cc56 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<string, number | string | undefined | boolean>;
+ urlParams?: Record<string, number | string | undefined | boolean | object>;
},
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 9568542f76..0695946fe0 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 989bd19cc3..ca331d1e34 100644
--- a/superset/sqllab/utils.py
+++ b/superset/sqllab/utils.py
@@ -23,7 +23,7 @@ import pyarrow as pa
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 597f961346..25158d25c8 100644
--- a/tests/integration_tests/sql_lab/api_tests.py
+++ b/tests/integration_tests/sql_lab/api_tests.py
@@ -57,7 +57,6 @@ class TestSqlLabApi(SupersetTestCase):
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 @@ class TestSqlLabApi(SupersetTestCase):
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 @@ class TestSqlLabApi(SupersetTestCase):
{"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 @@ class TestSqlLabApi(SupersetTestCase):
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 0dc4e26aca..30b8401cc6 100644
--- a/tests/integration_tests/sqllab_tests.py
+++ b/tests/integration_tests/sqllab_tests.py
@@ -461,6 +461,63 @@ class TestSqlLab(SupersetTestCase):
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