fix(sqllab): infinite running state on disconnect (#23669)

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

@@ -80,6 +80,7 @@ export const STOP_QUERY = 'STOP_QUERY';
 export const REQUEST_QUERY_RESULTS = 'REQUEST_QUERY_RESULTS';
 export const QUERY_SUCCESS = 'QUERY_SUCCESS';
 export const QUERY_FAILED = 'QUERY_FAILED';
+export const CLEAR_INACTIVE_QUERIES = 'CLEAR_INACTIVE_QUERIES';
 export const CLEAR_QUERY_RESULTS = 'CLEAR_QUERY_RESULTS';
 export const REMOVE_DATA_PREVIEW = 'REMOVE_DATA_PREVIEW';
 export const CHANGE_DATA_PREVIEW_ID = 'CHANGE_DATA_PREVIEW_ID';
@@ -219,6 +220,10 @@ export function estimateQueryCost(queryEditor) {
   };
 }
 
+export function clearInactiveQueries() {
+  return { type: CLEAR_INACTIVE_QUERIES };
+}
+
 export function startQuery(query) {
   Object.assign(query, {
     id: query.id ? query.id : shortid.generate(),

superset-frontend/src/SqlLab/components/App/index.jsx

@@ -168,7 +168,7 @@ class App extends React.PureComponent {
   }
 
   render() {
-    const { queries, actions, queriesLastUpdate } = this.props;
+    const { queries, queriesLastUpdate } = this.props;
     if (this.state.hash && this.state.hash === '#search') {
       return window.location.replace('/superset/sqllab/history/');
     }
@@ -176,7 +176,6 @@ class App extends React.PureComponent {
       <SqlLabStyles data-test="SqlLabApp" className="App SqlLab">
         <QueryAutoRefresh
           queries={queries}
-          refreshQueries={actions?.refreshQueries}
           queriesLastUpdate={queriesLastUpdate}
         />
         <TabbedSqlEditors />

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

@@ -16,15 +16,26 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+import fetchMock from 'fetch-mock';
 import React from 'react';
-import { render } from '@testing-library/react';
+import configureStore from 'redux-mock-store';
+import thunk from 'redux-thunk';
+import { render, waitFor } from 'spec/helpers/testing-library';
+import {
+  CLEAR_INACTIVE_QUERIES,
+  REFRESH_QUERIES,
+} from 'src/SqlLab/actions/sqlLab';
 import QueryAutoRefresh, {
   isQueryRunning,
   shouldCheckForQueries,
+  QUERY_UPDATE_FREQ,
 } from 'src/SqlLab/components/QueryAutoRefresh';
 import { successfulQuery, runningQuery } from 'src/SqlLab/fixtures';
 import { QueryDictionary } from 'src/SqlLab/types';
 
+const middlewares = [thunk];
+const mockStore = configureStore(middlewares);
+
 // NOTE: The uses of @ts-ignore in this file is to enable testing of bad inputs to verify the
 // function / component handles bad data elegantly
 describe('QueryAutoRefresh', () => {
@@ -34,10 +45,14 @@ describe('QueryAutoRefresh', () => {
   const successfulQueries: QueryDictionary = {};
   successfulQueries[successfulQuery.id] = successfulQuery;
 
-  const refreshQueries = jest.fn();
-
   const queriesLastUpdate = Date.now();
 
+  const refreshApi = 'glob:*/api/v1/query/updated_since?*';
+
+  afterEach(() => {
+    fetchMock.reset();
+  });
+
   it('isQueryRunning returns true for valid running query', () => {
     const running = isQueryRunning(runningQuery);
     expect(running).toBe(true);
@@ -91,43 +106,119 @@ describe('QueryAutoRefresh', () => {
     ).toBe(false);
   });
 
-  it('Attempts to refresh when given pending query', () => {
+  it('Attempts to refresh when given pending query', async () => {
+    const store = mockStore();
+    fetchMock.get(refreshApi, {
+      result: [
+        {
+          id: runningQuery.id,
+          status: 'success',
+        },
+      ],
+    });
+
     render(
       <QueryAutoRefresh
         queries={runningQueries}
-        refreshQueries={refreshQueries}
         queriesLastUpdate={queriesLastUpdate}
       />,
+      { useRedux: true, store },
+    );
+    await waitFor(
+      () =>
+        expect(store.getActions()).toContainEqual(
+          expect.objectContaining({
+            type: REFRESH_QUERIES,
+          }),
+        ),
+      { timeout: QUERY_UPDATE_FREQ + 100 },
     );
-    setTimeout(() => {
-      expect(refreshQueries).toHaveBeenCalled();
-    }, 1000);
   });
 
-  it('Does not fail and attempts to refresh when given pending query and invlaid query', () => {
+  it('Attempts to clear inactive queries when updated queries are empty', async () => {
+    const store = mockStore();
+    fetchMock.get(refreshApi, {
+      result: [],
+    });
+
+    render(
+      <QueryAutoRefresh
+        queries={runningQueries}
+        queriesLastUpdate={queriesLastUpdate}
+      />,
+      { useRedux: true, store },
+    );
+    await waitFor(
+      () =>
+        expect(store.getActions()).toContainEqual(
+          expect.objectContaining({
+            type: CLEAR_INACTIVE_QUERIES,
+          }),
+        ),
+      { timeout: QUERY_UPDATE_FREQ + 100 },
+    );
+    expect(
+      store.getActions().filter(({ type }) => type === REFRESH_QUERIES),
+    ).toHaveLength(0);
+    expect(fetchMock.calls(refreshApi)).toHaveLength(1);
+  });
+
+  it('Does not fail and attempts to refresh when given pending query and invlaid query', async () => {
+    const store = mockStore();
+    fetchMock.get(refreshApi, {
+      result: [
+        {
+          id: runningQuery.id,
+          status: 'success',
+        },
+      ],
+    });
+
     render(
       <QueryAutoRefresh
         // @ts-ignore
         queries={{ ...runningQueries, g324t: null }}
-        refreshQueries={refreshQueries}
         queriesLastUpdate={queriesLastUpdate}
       />,
+      { useRedux: true, store },
+    );
+    await waitFor(
+      () =>
+        expect(store.getActions()).toContainEqual(
+          expect.objectContaining({
+            type: REFRESH_QUERIES,
+          }),
+        ),
+      { timeout: QUERY_UPDATE_FREQ + 100 },
     );
-    setTimeout(() => {
-      expect(refreshQueries).toHaveBeenCalled();
-    }, 1000);
   });
 
-  it('Does NOT Attempt to refresh when given only completed queries', () => {
+  it('Does NOT Attempt to refresh when given only completed queries', async () => {
+    const store = mockStore();
+    fetchMock.get(refreshApi, {
+      result: [
+        {
+          id: runningQuery.id,
+          status: 'success',
+        },
+      ],
+    });
     render(
       <QueryAutoRefresh
         queries={successfulQueries}
-        refreshQueries={refreshQueries}
         queriesLastUpdate={queriesLastUpdate}
       />,
+      { useRedux: true, store },
     );
-    setTimeout(() => {
-      expect(refreshQueries).not.toHaveBeenCalled();
-    }, 1000);
+    await waitFor(
+      () =>
+        expect(store.getActions()).toContainEqual(
+          expect.objectContaining({
+            type: CLEAR_INACTIVE_QUERIES,
+          }),
+        ),
+      { timeout: QUERY_UPDATE_FREQ + 100 },
+    );
+    expect(fetchMock.calls(refreshApi)).toHaveLength(0);
   });
 });

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

@@ -16,7 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { useState } from 'react';
+import { useRef } from 'react';
+import { useDispatch } from 'react-redux';
 import { isObject } from 'lodash';
 import rison from 'rison';
 import {
@@ -27,19 +28,18 @@ import {
 } from '@superset-ui/core';
 import { QueryDictionary } from 'src/SqlLab/types';
 import useInterval from 'src/SqlLab/utils/useInterval';
+import {
+  refreshQueries,
+  clearInactiveQueries,
+} from 'src/SqlLab/actions/sqlLab';
 
-const QUERY_UPDATE_FREQ = 2000;
+export const QUERY_UPDATE_FREQ = 2000;
 const QUERY_UPDATE_BUFFER_MS = 5000;
 const MAX_QUERY_AGE_TO_POLL = 21600000;
 const QUERY_TIMEOUT_LIMIT = 10000;
 
-interface RefreshQueriesFunc {
-  (alteredQueries: any): any;
-}
-
 export interface QueryAutoRefreshProps {
   queries: QueryDictionary;
-  refreshQueries: RefreshQueriesFunc;
   queriesLastUpdate: number;
 }
 
@@ -61,20 +61,22 @@ export const shouldCheckForQueries = (queryList: QueryDictionary): boolean => {
 
 function QueryAutoRefresh({
   queries,
-  refreshQueries,
   queriesLastUpdate,
 }: QueryAutoRefreshProps) {
   // We do not want to spam requests in the case of slow connections and potentially receive responses out of order
   // pendingRequest check ensures we only have one active http call to check for query statuses
-  const [pendingRequest, setPendingRequest] = useState(false);
+  const pendingRequestRef = useRef(false);
+  const cleanInactiveRequestRef = useRef(false);
+  const dispatch = useDispatch();
 
   const checkForRefresh = () => {
-    if (!pendingRequest && shouldCheckForQueries(queries)) {
+    const shouldRequestChecking = shouldCheckForQueries(queries);
+    if (!pendingRequestRef.current && shouldRequestChecking) {
       const params = rison.encode({
         last_updated_ms: queriesLastUpdate - QUERY_UPDATE_BUFFER_MS,
       });
 
-      setPendingRequest(true);
+      pendingRequestRef.current = true;
       SupersetClient.get({
         endpoint: `/api/v1/query/updated_since?q=${params}`,
         timeout: QUERY_TIMEOUT_LIMIT,
@@ -82,19 +84,27 @@ function QueryAutoRefresh({
         .then(({ json }) => {
           if (json) {
             const jsonPayload = json as { result?: QueryResponse[] };
-            const queries =
-              jsonPayload?.result?.reduce((acc, current) => {
-                acc[current.id] = current;
-                return acc;
-              }, {}) ?? {};
-            refreshQueries?.(queries);
+            if (jsonPayload?.result?.length) {
+              const queries =
+                jsonPayload?.result?.reduce((acc, current) => {
+                  acc[current.id] = current;
+                  return acc;
+                }, {}) ?? {};
+              dispatch(refreshQueries(queries));
+            } else {
+              dispatch(clearInactiveQueries());
+            }
           }
         })
         .catch(() => {})
         .finally(() => {
-          setPendingRequest(false);
+          pendingRequestRef.current = false;
         });
     }
+    if (!cleanInactiveRequestRef.current && !shouldRequestChecking) {
+      dispatch(clearInactiveQueries());
+      cleanInactiveRequestRef.current = true;
+    }
   };
 
   // Solves issue where direct usage of setInterval in function components

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

@@ -742,6 +742,21 @@ export default function sqlLabReducer(state = {}, action) {
       }
       return { ...state, queries: newQueries, queriesLastUpdate };
     },
+    [actions.CLEAR_INACTIVE_QUERIES]() {
+      const { queries } = state;
+      const cleanedQueries = Object.fromEntries(
+        Object.entries(queries).filter(([, query]) => {
+          if (
+            ['running', 'pending'].includes(query.state) &&
+            query.progress === 0
+          ) {
+            return false;
+          }
+          return true;
+        }),
+      );
+      return { ...state, queries: cleanedQueries };
+    },
     [actions.SET_USER_OFFLINE]() {
       return { ...state, offline: action.offline };
     },