feat(sqllab): Add timeout on fetching query results (apache#29959)

(cherry picked from commit ff3b86b)

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

@@ -294,21 +294,25 @@ export function requestQueryResults(query) {
   return { type: REQUEST_QUERY_RESULTS, query };
 }
 
-export function fetchQueryResults(query, displayLimit) {
-  return function (dispatch) {
+export function fetchQueryResults(query, displayLimit, timeoutInMs) {
+  return function (dispatch, getState) {
+    const { SQLLAB_QUERY_RESULT_TIMEOUT } = getState().common?.conf ?? {};
     dispatch(requestQueryResults(query));
 
     const queryParams = rison.encode({
       key: query.resultsKey,
       rows: displayLimit || null,
     });
-
+    const timeout = timeoutInMs ?? SQLLAB_QUERY_RESULT_TIMEOUT;
+    const controller = new AbortController();
     return SupersetClient.get({
       endpoint: `/api/v1/sqllab/results/?q=${queryParams}`,
       parseMethod: 'json-bigint',
+      ...(timeout && { timeout, signal: controller.signal }),
     })
       .then(({ json }) => dispatch(querySuccess(query, json)))
-      .catch(response =>
+      .catch(response => {
+        controller.abort();
         getClientErrorObject(response).then(error => {
           const message =
             error.error ||
@@ -318,8 +322,8 @@ export function fetchQueryResults(query, displayLimit) {
           return dispatch(
             queryFailed(query, message, error.link, error.errors),
           );
-        }),
-      );
+        });
+      });
   };
 }
 

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

@@ -174,8 +174,9 @@ describe('async actions', () => {
 
   describe('fetchQueryResults', () => {
     const makeRequest = () => {
+      const store = mockStore(initialState);
       const request = actions.fetchQueryResults(query);
-      return request(dispatch);
+      return request(dispatch, store.getState);
     };
 
     it('makes the fetch request', () => {

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

@@ -32,6 +32,7 @@ import {
   initialState,
   user,
   queryWithNoQueryLimit,
+  failedQueryWithFrontendTimeoutErrors,
 } from 'src/SqlLab/fixtures';
 
 const mockedProps = {
@@ -104,6 +105,16 @@ const failedQueryWithErrorsState = {
     },
   },
 };
+const failedQueryWithTimeoutState = {
+  ...initialState,
+  sqlLab: {
+    ...initialState.sqlLab,
+    queries: {
+      [failedQueryWithFrontendTimeoutErrors.id]:
+        failedQueryWithFrontendTimeoutErrors,
+    },
+  },
+};
 
 const newProps = {
   displayLimit: 1001,
@@ -319,6 +330,18 @@ describe('ResultSet', () => {
     expect(screen.getByText('Database error')).toBeInTheDocument();
   });
 
+  test('should render a timeout error with a retrial button', async () => {
+    await waitFor(() => {
+      setup(
+        { ...mockedProps, queryId: failedQueryWithFrontendTimeoutErrors.id },
+        mockStore(failedQueryWithTimeoutState),
+      );
+    });
+    expect(
+      screen.getByRole('button', { name: /Retry fetching results/i }),
+    ).toBeInTheDocument();
+  });
+
   test('renders if there is no limit in query.results but has queryLimit', async () => {
     const query = {
       ...queries[0],

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

@@ -42,6 +42,7 @@ import {
   css,
   getNumberFormatter,
   getExtensionsRegistry,
+  ErrorTypeEnum,
 } from '@superset-ui/core';
 import ErrorMessageWithStackTrace from 'src/components/ErrorMessage/ErrorMessageWithStackTrace';
 import {
@@ -225,8 +226,8 @@ const ResultSet = ({
     reRunQueryIfSessionTimeoutErrorOnMount();
   }, [reRunQueryIfSessionTimeoutErrorOnMount]);
 
-  const fetchResults = (q: typeof query) => {
-    dispatch(fetchQueryResults(q, displayLimit));
+  const fetchResults = (q: typeof query, timeout?: number) => {
+    dispatch(fetchQueryResults(q, displayLimit, timeout));
   };
 
   const prevQuery = usePrevious(query);
@@ -549,7 +550,18 @@ const ResultSet = ({
           link={query.link}
           source="sqllab"
         />
-        {trackingUrl}
+        {(query?.extra?.errors?.[0] || query?.errors?.[0])?.error_type ===
+        ErrorTypeEnum.FRONTEND_TIMEOUT_ERROR ? (
+          <Button
+            className="sql-result-track-job"
+            buttonSize="small"
+            onClick={() => fetchResults(query, 0)}
+          >
+            {t('Retry fetching results')}
+          </Button>
+        ) : (
+          trackingUrl
+        )}
       </ResultlessStyles>
     );
   }

superset-frontend/src/SqlLab/fixtures.ts

@@ -22,6 +22,7 @@ import { ColumnKeyTypeType } from 'src/SqlLab/components/ColumnElement';
 import {
   DatasourceType,
   denormalizeTimestamp,
+  ErrorTypeEnum,
   GenericDataType,
   QueryResponse,
   QueryState,
@@ -546,6 +547,20 @@ export const failedQueryWithErrors = {
   tempTable: '',
 };
 
+export const failedQueryWithFrontendTimeoutErrors = {
+  ...failedQueryWithErrorMessage,
+  errors: [
+    {
+      error_type: ErrorTypeEnum.FRONTEND_TIMEOUT_ERROR,
+      message: 'Request timed out',
+      level: 'error',
+      extra: {
+        timeout: 10,
+      },
+    },
+  ],
+};
+
 const baseQuery: QueryResponse = {
   queryId: 567,
   dbId: 1,

superset/config.py

@@ -1047,6 +1047,10 @@ class CeleryConfig:  # pylint: disable=too-few-public-methods
 # timeout.
 SQLLAB_QUERY_COST_ESTIMATE_TIMEOUT = int(timedelta(seconds=10).total_seconds())
 
+# Timeout duration for SQL Lab fetching query results by the resultsKey.
+# 0 means no timeout.
+SQLLAB_QUERY_RESULT_TIMEOUT = 0
+
 # The cost returned by the databases is a relative value; in order to map the cost to
 # a tangible value you need to define a custom formatter that takes into consideration
 # your specific infrastructure. For example, you could analyze queries a posteriori by

superset/views/base.py

@@ -115,6 +115,7 @@
     "PREVENT_UNSAFE_DEFAULT_URLS_ON_DATASET",
     "JWT_ACCESS_CSRF_COOKIE_NAME",
     "SLACK_ENABLE_AVATARS",
+    "SQLLAB_QUERY_RESULT_TIMEOUT",
 )
 
 logger = logging.getLogger(__name__)