fix(oauth2): clear schema/catalog auth banner after OAuth2 redirect refetch (#41913)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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

@@ -18,8 +18,16 @@
  */
 import type { ReactChild } from 'react';
 import fetchMock from 'fetch-mock';
-import { render, screen, waitFor } from 'spec/helpers/testing-library';
+import {
+  act,
+  createStore,
+  render,
+  screen,
+  waitFor,
+} from 'spec/helpers/testing-library';
+import reducerIndex from 'spec/helpers/reducerIndex';
 import userEvent from '@testing-library/user-event';
+import { api } from 'src/hooks/apiResources/queryApi';
 import { initialState, defaultQueryEditor } from 'src/SqlLab/fixtures';
 
 import { ViewLocations } from 'src/SqlLab/contributions';
@@ -346,3 +354,64 @@ test('closes a schema while searchTerm is active and keeps it closed', async ()
   // The schema node itself remains visible as a matching ancestor (just collapsed)
   expect(screen.getByText('public')).toBeInTheDocument();
 });
+
+test('clears the OAuth error banner after a Tables invalidateTags refetch', async () => {
+  // Regression test for the OAuth2 crud symptom (follow-up to PR #41101).
+  // Expanding a schema lazily fetches its tables; when that fails with an
+  // OAuth2 auth error the banner is held in local reducer state and, before
+  // this fix, only cleared when a table list loaded via manual re-expansion.
+  // After the OAuth2 redirect, OAuth2RedirectMessage dispatches
+  // invalidateTags(['Tables']); the errored node's subscribed tables query must
+  // now refetch automatically and clear the banner.
+  fetchMock.removeRoutes().clearHistory();
+  fetchMock.get('glob:*/api/v1/database/1/schemas/?*', {
+    count: mockSchemas.length,
+    result: mockSchemas,
+  });
+  let tablesShouldFail = true;
+  fetchMock.get('glob:*/api/v1/database/1/tables/*', () =>
+    tablesShouldFail
+      ? {
+          status: 500,
+          body: {
+            errors: [
+              {
+                error_type: 'GENERIC_DB_ENGINE_ERROR',
+                level: 'error',
+                message: 'Tables could not be loaded',
+                extra: {},
+              },
+            ],
+          },
+        }
+      : { count: mockTables.length, result: mockTables },
+  );
+
+  const store = createStore(getInitialState(), reducerIndex);
+  render(<TableExploreTree queryEditorId={mockedQueryEditorId} />, {
+    useRedux: true,
+    store,
+  });
+
+  await waitFor(() => {
+    expect(screen.getByText('public')).toBeInTheDocument();
+  });
+
+  // Expand the schema node → its table fetch rejects → the banner appears.
+  await userEvent.click(screen.getByText('public'));
+  expect(await screen.findByText('Unexpected error')).toBeInTheDocument();
+
+  // The OAuth2 redirect completes: the stored token makes the next fetch
+  // succeed, and the redirect handler invalidates the Tables cache.
+  tablesShouldFail = false;
+  act(() => {
+    store.dispatch(api.util.invalidateTags(['Tables']));
+  });
+
+  // The subscribed tables query refetches, the tables load, and the banner is
+  // cleared without the user manually re-expanding the node.
+  expect(await screen.findByText('users')).toBeInTheDocument();
+  await waitFor(() =>
+    expect(screen.queryByText('Unexpected error')).not.toBeInTheDocument(),
+  );
+});

superset-frontend/src/SqlLab/components/TableExploreTree/useTreeData.ts

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { useMemo, useReducer, useCallback } from 'react';
+import { useMemo, useReducer, useCallback, useEffect, useRef } from 'react';
 import { useAppDispatch } from 'src/SqlLab/hooks/useAppDispatch';
 import { t } from '@apache-superset/core/translation';
 import {
@@ -24,21 +24,33 @@ import {
   type TableMetaData,
   useSchemas,
   useLazyTablesQuery,
+  useTablesQuery,
   useLazyTableMetadataQuery,
   useLazyTableExtendedMetadataQuery,
 } from 'src/hooks/apiResources';
 import { addDangerToast } from 'src/SqlLab/actions/sqlLab';
 import type { TreeNodeData } from './types';
-import { SupersetError } from '@superset-ui/core';
+import { ClientErrorObject, SupersetError } from '@superset-ui/core';
 
 export const EMPTY_NODE_ID_PREFIX = 'empty:';
 
+// Identifies the schema node whose table list failed to load, so the tree can
+// automatically recover once the underlying Tables cache is refetched (e.g.
+// after the OAuth2 redirect dispatches invalidateTags(['Tables'])).
+interface ErroredNode {
+  schemaKey: string;
+  dbId: number;
+  catalog: string | null | undefined;
+  schema: string;
+}
+
 // Reducer state and actions
 interface TreeDataState {
   tableData: Record<string, { options: Table[] }>;
   tableSchemaData: Record<string, TableMetaData>;
   loadingNodes: Record<string, boolean>;
   errorPayload: SupersetError | null;
+  erroredNode: ErroredNode | null;
 }
 
 type TreeDataAction =
@@ -46,13 +58,18 @@ type TreeDataAction =
   | { type: 'SET_TABLE_SCHEMA_DATA'; key: string; data: TableMetaData }
   | { type: 'CLEAR_TABLE_SCHEMA_DATA'; key: string }
   | { type: 'SET_LOADING_NODE'; nodeId: string; loading: boolean }
-  | { type: 'SET_ERROR'; errorPayload: SupersetError | null };
+  | {
+      type: 'SET_ERROR';
+      errorPayload: SupersetError | null;
+      erroredNode: ErroredNode | null;
+    };
 
 const initialState: TreeDataState = {
   tableData: {},
   tableSchemaData: {},
   loadingNodes: {},
   errorPayload: null,
+  erroredNode: null,
 };
 
 function treeDataReducer(
@@ -64,6 +81,7 @@ function treeDataReducer(
       return {
         ...state,
         errorPayload: null,
+        erroredNode: null,
         tableData: { ...state.tableData, [action.key]: action.data },
       };
     case 'SET_TABLE_SCHEMA_DATA':
@@ -90,6 +108,7 @@ function treeDataReducer(
       return {
         ...state,
         errorPayload: action.errorPayload,
+        erroredNode: action.erroredNode,
       };
 
     default:
@@ -144,7 +163,67 @@ const useTreeData = ({
 
   // Combined state for table data, schema data, loading nodes, and data version
   const [state, dispatch] = useReducer(treeDataReducer, initialState);
-  const { tableData, tableSchemaData, loadingNodes, errorPayload } = state;
+  const {
+    tableData,
+    tableSchemaData,
+    loadingNodes,
+    errorPayload,
+    erroredNode,
+  } = state;
+
+  // Tables are loaded lazily on node toggle, so a schema whose table list fails
+  // (e.g. an OAuth2 auth error) has no active subscription that would recover on
+  // cache invalidation. Subscribe to the tables query for the single errored
+  // node so that when OAuth2RedirectMessage dispatches invalidateTags(['Tables'])
+  // after the redirect, this query refetches automatically. The subscribed
+  // entry shares the cache key (dbId + schema) that the lazy fetch already
+  // populated with the error, so this reflects that error and does not trigger
+  // an eager refetch of its own.
+  const erroredTablesResult = useTablesQuery(
+    {
+      dbId: erroredNode?.dbId,
+      catalog: erroredNode?.catalog,
+      schema: erroredNode?.schema,
+      forceRefresh: false,
+    },
+    { skip: !erroredNode },
+  );
+  const wasFetchingErroredRef = useRef(false);
+
+  useEffect(() => {
+    // Recover the errored schema node when its subscribed tables query finishes
+    // a fetch (driven by the Tables cache invalidation). Keying off the
+    // isFetching true->false transition avoids acting on the initial rejected
+    // state and on unrelated re-renders. On success, SET_TABLE_DATA repopulates
+    // the node and clears the banner; on renewed failure the banner is re-armed.
+    if (!erroredNode) {
+      wasFetchingErroredRef.current = erroredTablesResult.isFetching;
+      return;
+    }
+    const { isSuccess, isError, isFetching, currentData, error } =
+      erroredTablesResult;
+    const nodeId = `schema:${erroredNode.dbId}:${erroredNode.schema}`;
+    if (isFetching && !wasFetchingErroredRef.current) {
+      dispatch({ type: 'SET_LOADING_NODE', nodeId, loading: true });
+    }
+    if (!isFetching && wasFetchingErroredRef.current) {
+      if (isSuccess && currentData) {
+        dispatch({
+          type: 'SET_TABLE_DATA',
+          key: erroredNode.schemaKey,
+          data: currentData,
+        });
+      } else if (isError) {
+        dispatch({
+          type: 'SET_ERROR',
+          errorPayload: (error as ClientErrorObject)?.errors?.[0] ?? null,
+          erroredNode,
+        });
+      }
+      dispatch({ type: 'SET_LOADING_NODE', nodeId, loading: false });
+    }
+    wasFetchingErroredRef.current = isFetching;
+  }, [erroredTablesResult, erroredNode]);
 
   // Shared helper: fetch table metadata + extended metadata and store in state.
   // preferCacheValue=true on initial open (use cached data if available),
@@ -233,6 +312,12 @@ const useTreeData = ({
               dispatch({
                 type: 'SET_ERROR',
                 errorPayload: error?.errors?.[0] ?? null,
+                erroredNode: {
+                  schemaKey,
+                  dbId: parsedDbId,
+                  catalog,
+                  schema,
+                },
               });
             })
             .finally(() => {
@@ -296,6 +381,12 @@ const useTreeData = ({
           dispatch({
             type: 'SET_ERROR',
             errorPayload: error?.errors?.[0] ?? null,
+            erroredNode: {
+              schemaKey,
+              dbId: refreshDbId,
+              catalog: refreshCatalog,
+              schema,
+            },
           });
         })
         .finally(() => {