pt
/
superset
mirror of https://github.com/apache/superset.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

perf(sqllab): reduce bootstrap data delay by queries (#27488)

This commit is contained in:
JUST.in DO IT 2024-03-18 12:52:23 -07:00 committed by GitHub
parent 376bfd05bd
commit f4bdcb5743
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 577 additions and 76 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
superset-frontend/src/SqlLab/components/QueryHistory/QueryHistory.test.tsx
View File

@ -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,
const emptyStateText = screen.getByText(
/run a query to display query history/i,
);
expect(emptyStateText).toBeVisible();
});
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();
});

106
superset-frontend/src/SqlLab/components/QueryHistory/index.tsx
View File

@ -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

21
superset-frontend/src/SqlLab/reducers/getInitialState.test.ts
View File

@ -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),

7
superset-frontend/src/SqlLab/reducers/getInitialState.ts
View File

@ -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

154
superset-frontend/src/hooks/apiResources/queries.test.ts Normal file
View File

@ -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),
],
});
});

176
superset-frontend/src/hooks/apiResources/queries.ts Normal file
View File

@ -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;

3
superset-frontend/src/hooks/apiResources/queryApi.ts
View File

@ -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,

19
superset/queries/api.py
View File

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

15
superset/sqllab/utils.py
View File

@ -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,
}

22
tests/integration_tests/sql_lab/api_tests.py
View File

@ -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",

57
tests/integration_tests/sqllab_tests.py
View File

@ -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