Skip to content
This repository has been archived by the owner on Nov 3, 2023. It is now read-only.

Commit

Permalink
feat(trino): support early cancellation of queries (apache#22498)
Browse files Browse the repository at this point in the history
(cherry picked from commit b6d39d1)
  • Loading branch information
villebro authored and john-bodley committed Jan 25, 2023
1 parent 92fcf9b commit c2715d6
Show file tree
Hide file tree
Showing 16 changed files with 230 additions and 73 deletions.
6 changes: 4 additions & 2 deletions superset-frontend/src/SqlLab/actions/sqlLab.js
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
* under the License.
*/
import shortid from 'shortid';
import { t, SupersetClient } from '@superset-ui/core';
import { SupersetClient, t } from '@superset-ui/core';
import invert from 'lodash/invert';
import mapKeys from 'lodash/mapKeys';
import { isFeatureEnabled, FeatureFlag } from 'src/featureFlags';
Expand Down Expand Up @@ -229,11 +229,13 @@ export function startQuery(query) {

export function querySuccess(query, results) {
return function (dispatch) {
const sqlEditorId = results?.query?.sqlEditorId;
const sync =
sqlEditorId &&
!query.isDataPreview &&
isFeatureEnabled(FeatureFlag.SQLLAB_BACKEND_PERSISTENCE)
? SupersetClient.put({
endpoint: encodeURI(`/tabstateview/${results.query.sqlEditorId}`),
endpoint: encodeURI(`/tabstateview/${sqlEditorId}`),
postPayload: { latest_query_id: query.id },
})
: Promise.resolve();
Expand Down
31 changes: 31 additions & 0 deletions superset-frontend/src/SqlLab/actions/sqlLab.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ import {
initialState,
queryId,
} from 'src/SqlLab/fixtures';
import { QueryState } from '@superset-ui/core';

const middlewares = [thunk];
const mockStore = configureMockStore(middlewares);
Expand Down Expand Up @@ -502,6 +503,7 @@ describe('async actions', () => {
const results = {
data: mockBigNumber,
query: { sqlEditorId: 'abcd' },
status: QueryState.SUCCESS,
query_id: 'efgh',
};
fetchMock.get(fetchQueryEndpoint, JSON.stringify(results), {
Expand All @@ -525,6 +527,35 @@ describe('async actions', () => {
expect(fetchMock.calls(updateTabStateEndpoint)).toHaveLength(1);
});
});

it("doesn't update the tab state in the backend on stoppped query", () => {
expect.assertions(2);

const results = {
status: QueryState.STOPPED,
query_id: 'efgh',
};
fetchMock.get(fetchQueryEndpoint, JSON.stringify(results), {
overwriteRoutes: true,
});
const store = mockStore({});
const expectedActions = [
{
type: actions.REQUEST_QUERY_RESULTS,
query,
},
// missing below
{
type: actions.QUERY_SUCCESS,
query,
results,
},
];
return store.dispatch(actions.fetchQueryResults(query)).then(() => {
expect(store.getActions()).toEqual(expectedActions);
expect(fetchMock.calls(updateTabStateEndpoint)).toHaveLength(0);
});
});
});

describe('addQueryEditor', () => {
Expand Down
26 changes: 14 additions & 12 deletions superset-frontend/src/SqlLab/components/ResultSet/index.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,13 @@
* specific language governing permissions and limitations
* under the License.
*/
import React, { useState, useEffect, useCallback } from 'react';
import React, { useCallback, useEffect, useState } from 'react';
import { useDispatch } from 'react-redux';
import ButtonGroup from 'src/components/ButtonGroup';
import Alert from 'src/components/Alert';
import Button from 'src/components/Button';
import shortid from 'shortid';
import { styled, t, QueryResponse } from '@superset-ui/core';
import { QueryResponse, QueryState, styled, t } from '@superset-ui/core';
import { usePrevious } from 'src/hooks/usePrevious';
import ErrorMessageWithStackTrace from 'src/components/ErrorMessage/ErrorMessageWithStackTrace';
import {
Expand All @@ -43,9 +43,9 @@ import CopyToClipboard from 'src/components/CopyToClipboard';
import { addDangerToast } from 'src/components/MessageToasts/actions';
import { prepareCopyToClipboardTabularData } from 'src/utils/common';
import {
CtasEnum,
clearQueryResults,
addQueryEditor,
clearQueryResults,
CtasEnum,
fetchQueryResults,
reFetchQueryResults,
reRunQuery,
Expand Down Expand Up @@ -387,8 +387,8 @@ const ResultSet = ({
let trackingUrl;
if (
query.trackingUrl &&
query.state !== 'success' &&
query.state !== 'fetching'
query.state !== QueryState.SUCCESS &&
query.state !== QueryState.FETCHING
) {
trackingUrl = (
<Button
Expand All @@ -397,7 +397,9 @@ const ResultSet = ({
href={query.trackingUrl}
target="_blank"
>
{query.state === 'running' ? t('Track job') : t('See query details')}
{query.state === QueryState.RUNNING
? t('Track job')
: t('See query details')}
</Button>
);
}
Expand All @@ -406,11 +408,11 @@ const ResultSet = ({
sql = <HighlightedSql sql={query.sql} />;
}

if (query.state === 'stopped') {
if (query.state === QueryState.STOPPED) {
return <Alert type="warning" message={t('Query was stopped')} />;
}

if (query.state === 'failed') {
if (query.state === QueryState.FAILED) {
return (
<ResultlessStyles>
<ErrorMessageWithStackTrace
Expand All @@ -426,7 +428,7 @@ const ResultSet = ({
);
}

if (query.state === 'success' && query.ctas) {
if (query.state === QueryState.SUCCESS && query.ctas) {
const { tempSchema, tempTable } = query;
let object = 'Table';
if (query.ctas_method === CtasEnum.VIEW) {
Expand Down Expand Up @@ -465,7 +467,7 @@ const ResultSet = ({
);
}

if (query.state === 'success' && query.results) {
if (query.state === QueryState.SUCCESS && query.results) {
const { results } = query;
// Accounts for offset needed for height of ResultSetRowsReturned component if !limitReached
const rowMessageHeight = !limitReached ? 32 : 0;
Expand Down Expand Up @@ -508,7 +510,7 @@ const ResultSet = ({
}
}

if (query.cached || (query.state === 'success' && !query.results)) {
if (query.cached || (query.state === QueryState.SUCCESS && !query.results)) {
if (query.isDataPreview) {
return (
<Button
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ const SqlEditorTabHeader: React.FC<Props> = ({ queryEditor }) => {
}),
shallowEqual,
);
const queryStatus = useSelector<SqlLabRootState, QueryState>(
const queryState = useSelector<SqlLabRootState, QueryState>(
({ sqlLab }) => sqlLab.queries[qe.latestQueryId || '']?.state || '',
);
const dispatch = useDispatch();
Expand Down Expand Up @@ -139,7 +139,7 @@ const SqlEditorTabHeader: React.FC<Props> = ({ queryEditor }) => {
</Menu>
}
/>
<TabTitle>{qe.name}</TabTitle> <TabStatusIcon tabState={queryStatus} />{' '}
<TabTitle>{qe.name}</TabTitle> <TabStatusIcon tabState={queryState} />{' '}
</TabTitleWrapper>
);
};
Expand Down
29 changes: 16 additions & 13 deletions superset-frontend/src/SqlLab/reducers/sqlLab.js
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,7 @@
* specific language governing permissions and limitations
* under the License.
*/
import { t } from '@superset-ui/core';

import { QueryState, t } from '@superset-ui/core';
import getInitialState from './getInitialState';
import * as actions from '../actions/sqlLab';
import { now } from '../../utils/dates';
Expand Down Expand Up @@ -391,7 +390,7 @@ export default function sqlLabReducer(state = {}, action) {
},
[actions.STOP_QUERY]() {
return alterInObject(state, 'queries', action.query, {
state: 'stopped',
state: QueryState.STOPPED,
results: [],
});
},
Expand All @@ -405,20 +404,24 @@ export default function sqlLabReducer(state = {}, action) {
},
[actions.REQUEST_QUERY_RESULTS]() {
return alterInObject(state, 'queries', action.query, {
state: 'fetching',
state: QueryState.FETCHING,
});
},
[actions.QUERY_SUCCESS]() {
// prevent race condition were query succeeds shortly after being canceled
if (action.query.state === 'stopped') {
// prevent race condition where query succeeds shortly after being canceled
// or the final result was unsuccessful
if (
action.query.state === QueryState.STOPPED ||
action.results.status !== QueryState.SUCCESS
) {
return state;
}
const alts = {
endDttm: now(),
progress: 100,
results: action.results,
rows: action?.results?.query?.rows || 0,
state: 'success',
state: QueryState.SUCCESS,
limitingFactor: action?.results?.query?.limitingFactor,
tempSchema: action?.results?.query?.tempSchema,
tempTable: action?.results?.query?.tempTable,
Expand All @@ -434,11 +437,11 @@ export default function sqlLabReducer(state = {}, action) {
return alterInObject(state, 'queries', action.query, alts);
},
[actions.QUERY_FAILED]() {
if (action.query.state === 'stopped') {
if (action.query.state === QueryState.STOPPED) {
return state;
}
const alts = {
state: 'failed',
state: QueryState.FAILED,
errors: action.errors,
errorMessage: action.msg,
endDttm: now(),
Expand Down Expand Up @@ -723,8 +726,8 @@ export default function sqlLabReducer(state = {}, action) {
Object.entries(action.alteredQueries).forEach(([id, changedQuery]) => {
if (
!state.queries.hasOwnProperty(id) ||
(state.queries[id].state !== 'stopped' &&
state.queries[id].state !== 'failed')
(state.queries[id].state !== QueryState.STOPPED &&
state.queries[id].state !== QueryState.FAILED)
) {
if (changedQuery.changedOn > queriesLastUpdate) {
queriesLastUpdate = changedQuery.changedOn;
Expand All @@ -738,8 +741,8 @@ export default function sqlLabReducer(state = {}, action) {
// because of async behavior, sql lab may still poll a couple of seconds
// when it started fetching or finished rendering results
state:
currentState === 'success' &&
['fetching', 'success'].includes(prevState)
currentState === QueryState.SUCCESS &&
[QueryState.FETCHING, QueryState.SUCCESS].includes(prevState)
? prevState
: currentState,
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import ListView from 'src/components/ListView';
import Filters from 'src/components/ListView/Filters';
import SyntaxHighlighter from 'react-syntax-highlighter/dist/cjs/light';
import SubMenu from 'src/views/components/SubMenu';
import { QueryState } from '@superset-ui/core';

// store needed for withToasts
const mockStore = configureStore([thunk]);
Expand All @@ -54,7 +55,7 @@ const mockQueries: QueryObject[] = [...new Array(3)].map((_, i) => ({
{ schema: 'foo', table: 'table' },
{ schema: 'bar', table: 'table_2' },
],
status: 'success',
status: QueryState.SUCCESS,
tab_name: 'Main Tab',
user: {
first_name: 'cool',
Expand Down
34 changes: 26 additions & 8 deletions superset-frontend/src/views/CRUD/data/query/QueryList.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,13 @@
* under the License.
*/
import React, { useMemo, useState, useCallback, ReactElement } from 'react';
import { SupersetClient, t, styled, useTheme } from '@superset-ui/core';
import {
QueryState,
styled,
SupersetClient,
t,
useTheme,
} from '@superset-ui/core';
import moment from 'moment';
import {
createFetchRelated,
Expand Down Expand Up @@ -127,41 +133,53 @@ function QueryList({ addDangerToast }: QueryListProps) {
row: {
original: { status },
},
}: any) => {
}: {
row: {
original: {
status: QueryState;
};
};
}) => {
const statusConfig: {
name: ReactElement | null;
label: string;
} = {
name: null,
label: '',
};
if (status === 'success') {
if (status === QueryState.SUCCESS) {
statusConfig.name = (
<Icons.Check iconColor={theme.colors.success.base} />
);
statusConfig.label = t('Success');
} else if (status === 'failed' || status === 'stopped') {
} else if (
status === QueryState.FAILED ||
status === QueryState.STOPPED
) {
statusConfig.name = (
<Icons.XSmall
iconColor={
status === 'failed'
status === QueryState.FAILED
? theme.colors.error.base
: theme.colors.grayscale.base
}
/>
);
statusConfig.label = t('Failed');
} else if (status === 'running') {
} else if (status === QueryState.RUNNING) {
statusConfig.name = (
<Icons.Running iconColor={theme.colors.primary.base} />
);
statusConfig.label = t('Running');
} else if (status === 'timed_out') {
} else if (status === QueryState.TIMED_OUT) {
statusConfig.name = (
<Icons.Offline iconColor={theme.colors.grayscale.light1} />
);
statusConfig.label = t('Offline');
} else if (status === 'scheduled' || status === 'pending') {
} else if (
status === QueryState.SCHEDULED ||
status === QueryState.PENDING
) {
statusConfig.name = (
<Icons.Queued iconColor={theme.colors.grayscale.base} />
);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import QueryPreviewModal from 'src/views/CRUD/data/query/QueryPreviewModal';
import { QueryObject } from 'src/views/CRUD/types';
import SyntaxHighlighter from 'react-syntax-highlighter/dist/cjs/light';
import { act } from 'react-dom/test-utils';
import { QueryState } from '@superset-ui/core';

// store needed for withToasts
const mockStore = configureStore([thunk]);
Expand All @@ -46,7 +47,7 @@ const mockQueries: QueryObject[] = [...new Array(3)].map((_, i) => ({
{ schema: 'foo', table: 'table' },
{ schema: 'bar', table: 'table_2' },
],
status: 'success',
status: QueryState.SUCCESS,
tab_name: 'Main Tab',
user: {
first_name: 'cool',
Expand Down
Loading

0 comments on commit c2715d6

Please sign in to comment.