Move selectedTable state up into ResultsApp

Author: asgerf

extensions/ql-vscode/src/view/results/ResultTables.tsx

@@ -13,7 +13,6 @@ import {
   ALERTS_TABLE_NAME,
   GRAPH_TABLE_NAME,
   SELECT_TABLE_NAME,
-  getDefaultResultSetName,
 } from "../../common/interface-types";
 import { tableHeaderClassName } from "./result-table-utils";
 import { vscode } from "../vscode-api";
@@ -42,6 +41,8 @@ interface ResultTablesProps {
   isLoadingNewResults: boolean;
   queryName: string;
   queryPath: string;
+  selectedTable: string;
+  onSelectedTableChange: (tableName: string) => void;
   problemsViewSelected: boolean;
   onProblemsViewSelectedChange: (selected: boolean) => void;
 }
@@ -102,30 +103,12 @@ export function ResultTables(props: ResultTablesProps) {
     origResultsPaths,
     isLoadingNewResults,
     sortStates,
+    selectedTable,
+    onSelectedTableChange,
     problemsViewSelected,
     onProblemsViewSelectedChange,
   } = props;
 
-  const [selectedTable, setSelectedTable] = useState(
-    parsedResultSets.selectedTable ||
-      getDefaultResultSet(getResultSets(rawResultSets, interpretation)),
-  );
-  useEffect(() => {
-    const resultSetExists =
-      parsedResultSets.resultSetNames.some((v) => selectedTable === v) ||
-      getResultSets(rawResultSets, interpretation).some(
-        (v) => selectedTable === getResultSetName(v),
-      );
-
-    // If the selected result set does not exist, select the default result set.
-    if (!resultSetExists) {
-      setSelectedTable(
-        parsedResultSets.selectedTable ||
-          getDefaultResultSet(getResultSets(rawResultSets, interpretation)),
-      );
-    }
-  }, [parsedResultSets, interpretation, rawResultSets, selectedTable]);
-
   const onTableSelectionChange = useCallback(
     (event: React.ChangeEvent<HTMLSelectElement>): void => {
       const selectedTable = event.target.value;
@@ -134,9 +117,10 @@ export function ResultTables(props: ResultTablesProps) {
         pageNumber: 0,
         selectedTable,
       });
+      onSelectedTableChange(selectedTable);
       sendTelemetry("local-results-table-selection");
     },
-    [],
+    [onSelectedTableChange],
   );
 
   const handleCheckboxChanged = useCallback(
@@ -227,7 +211,7 @@ export function ResultTables(props: ResultTablesProps) {
           sortState={sortStates.get(resultSetName)}
           nonemptyRawResults={nonemptyRawResults}
           showRawResults={() => {
-            setSelectedTable(SELECT_TABLE_NAME);
+            onSelectedTableChange(SELECT_TABLE_NAME);
             sendTelemetry("local-results-show-raw-results");
           }}
           offset={offset}
@@ -237,12 +221,6 @@ export function ResultTables(props: ResultTablesProps) {
   );
 }
 
-function getDefaultResultSet(resultSets: readonly ResultSet[]): string {
-  return getDefaultResultSetName(
-    resultSets.map((resultSet) => getResultSetName(resultSet)),
-  );
-}
-
 function getResultSetName(resultSet: ResultSet): string {
   switch (resultSet.t) {
     case "RawResultSet":

extensions/ql-vscode/src/view/results/ResultsApp.tsx

@@ -15,6 +15,7 @@ import {
   ALERTS_TABLE_NAME,
   DEFAULT_USER_SETTINGS,
   GRAPH_TABLE_NAME,
+  getDefaultResultSetName,
 } from "../../common/interface-types";
 import { useMessageFromExtension } from "../common/useMessageFromExtension";
 import { ResultTables } from "./ResultTables";
@@ -64,6 +65,7 @@ interface ResultsViewState {
   displayedResults: ResultsState;
   nextResultsInfo: ResultsInfo | null;
   isExpectingResultsUpdate: boolean;
+  selectedTable: string | undefined;
 }
 
 /**
@@ -78,6 +80,7 @@ export function ResultsApp() {
     },
     nextResultsInfo: null,
     isExpectingResultsUpdate: true,
+    selectedTable: undefined,
   });
 
   const [userSettings, setUserSettings] = useState<UserSettings>(
@@ -86,6 +89,39 @@ export function ResultsApp() {
 
   const [problemsViewSelected, setProblemsViewSelected] = useState(false);
 
+  const onSelectedTableChange = useCallback((tableName: string) => {
+    setState((prev) => {
+      if (tableName === prev.selectedTable) return prev;
+      return {
+        ...prev,
+        selectedTable: tableName,
+      };
+    });
+  }, []);
+
+  // Ensure selectedTable is valid for the current result sets.
+  // This runs in ResultsApp (not ResultTables) so it survives remounts.
+  const displayedResultsInfo = state.displayedResults.resultsInfo;
+  useEffect(() => {
+    if (!displayedResultsInfo) return;
+    const { parsedResultSets, interpretation } = displayedResultsInfo;
+    const allNames = interpretation
+      ? parsedResultSets.resultSetNames.concat([
+          interpretation.data.t === "GraphInterpretationData"
+            ? GRAPH_TABLE_NAME
+            : ALERTS_TABLE_NAME,
+        ])
+      : parsedResultSets.resultSetNames;
+    if (
+      state.selectedTable === undefined ||
+      !allNames.includes(state.selectedTable)
+    ) {
+      const tableName =
+        parsedResultSets.selectedTable ?? getDefaultResultSetName(allNames);
+      onSelectedTableChange(tableName);
+    }
+  }, [displayedResultsInfo, state.selectedTable, onSelectedTableChange]);
+
   const updateStateWithNewResultsInfo = useCallback(
     (resultsInfo: ResultsInfo): void => {
       let results: Results | null = null;
@@ -103,15 +139,16 @@ export function ResultsApp() {
         statusText = `Error loading results: ${errorMessage}`;
       }
 
-      setState({
+      setState((prev) => ({
+        ...prev,
         displayedResults: {
           resultsInfo,
           results,
           errorMessage: statusText,
         },
         nextResultsInfo: null,
         isExpectingResultsUpdate: false,
-      });
+      }));
     },
     [],
   );
@@ -232,6 +269,8 @@ export function ResultsApp() {
         }
         queryName={displayedResults.resultsInfo.queryName}
         queryPath={displayedResults.resultsInfo.queryPath}
+        selectedTable={state.selectedTable ?? ""}
+        onSelectedTableChange={onSelectedTableChange}
         problemsViewSelected={problemsViewSelected}
         onProblemsViewSelectedChange={setProblemsViewSelected}
       />