feat(trino): support early cancellation of queries (apache#22498)

(cherry picked from commit b6d39d1)

superset-frontend/src/SqlLab/actions/sqlLab.js

@@ -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';
@@ -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();

superset-frontend/src/SqlLab/actions/sqlLab.test.js

@@ -30,6 +30,7 @@ import {
   initialState,
   queryId,
 } from 'src/SqlLab/fixtures';
+import { QueryState } from '@superset-ui/core';
 
 const middlewares = [thunk];
 const mockStore = configureMockStore(middlewares);
@@ -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), {
@@ -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', () => {

superset-frontend/src/SqlLab/components/ResultSet/index.tsx

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

superset-frontend/src/SqlLab/components/SqlEditorTabHeader/index.tsx

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

superset-frontend/src/SqlLab/reducers/sqlLab.js

@@ -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';
@@ -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: [],
       });
     },
@@ -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,
@@ -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(),
@@ -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;
@@ -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,
           };

superset-frontend/src/views/CRUD/data/query/QueryList.test.tsx

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

superset-frontend/src/views/CRUD/data/query/QueryList.tsx

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

superset-frontend/src/views/CRUD/data/query/QueryPreviewModal.test.tsx

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