Improve auto-modeling in-progress state handling

extensions/ql-vscode/src/common/interface-types.ts

@@ -512,6 +512,7 @@ interface AddModeledMethodsMessage {
 
 interface SetInProgressMethodsMessage {
   t: "setInProgressMethods";
+  packageName: string;
   inProgressMethods: string[];
 }
 

extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts

@@ -32,6 +32,7 @@ export class AutoModeler {
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly setInProgressMethods: (
+      packageName: string,
       inProgressMethods: string[],
     ) => Promise<void>,
     private readonly addModeledMethods: (
@@ -120,6 +121,7 @@ export class AutoModeler {
           const candidatesToProcess = allCandidateMethods.slice(start, end);
 
           await this.setInProgressMethods(
+            packageName,
             candidatesToProcess.map((c) => c.signature),
           );
 
@@ -141,7 +143,7 @@ export class AutoModeler {
         }
       } finally {
         // Clear out in progress methods
-        await this.setInProgressMethods([]);
+        await this.setInProgressMethods(packageName, []);
       }
     });
   }

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -83,9 +83,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
       queryRunner,
       queryStorageDir,
       databaseItem,
-      async (inProgressMethods) => {
+      async (packageName, inProgressMethods) => {
         await this.postMessage({
           t: "setInProgressMethods",
+          packageName,
           inProgressMethods,
         });
       },

extensions/ql-vscode/src/data-extensions-editor/shared/in-progress-methods.ts

@@ -0,0 +1,31 @@
+/**
+ * A class that keeps track of which methods are in progress for each package.
+ */
+export class InProgressMethods {
+  // A map of in-progress method signatures for each package.
+  private readonly methodMap: Map<string, Set<string>>;
+
+  constructor() {
+    this.methodMap = new Map<string, Set<string>>();
+  }
+
+  public setPackageMethods(packageName: string, methods: Set<string>): void {
+    this.methodMap.set(packageName, methods);
+  }
+
+  public hasMethod(packageName: string, method: string): boolean {
+    const methods = this.methodMap.get(packageName);
+    if (methods) {
+      return methods.has(method);
+    }
+    return false;
+  }
+
+  public static fromExisting(methods: InProgressMethods): InProgressMethods {
+    const newInProgressMethods = new InProgressMethods();
+    methods.methodMap.forEach((value, key) => {
+      newInProgressMethods.methodMap.set(key, new Set<string>(value));
+    });
+    return newInProgressMethods;
+  }
+}

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -17,6 +17,7 @@ import { DataExtensionEditorViewState } from "../../data-extensions-editor/share
 import { ModeledMethodsList } from "./ModeledMethodsList";
 import { percentFormatter } from "./formatters";
 import { Mode } from "../../data-extensions-editor/shared/mode";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 import { getLanguageDisplayName } from "../../common/query-language";
 
 const LoadingContainer = styled.div`
@@ -91,8 +92,8 @@ export function DataExtensionsEditor({
     new Set(),
   );
 
-  const [inProgressSignatures, setInProgressSignatures] = useState<Set<string>>(
-    new Set(),
+  const [inProgressMethods, setInProgressMethods] = useState<InProgressMethods>(
+    new InProgressMethods(),
   );
 
   const [hideModeledApis, setHideModeledApis] = useState(true);
@@ -140,7 +141,15 @@ export function DataExtensionsEditor({
             );
             break;
           case "setInProgressMethods":
-            setInProgressSignatures(new Set(msg.inProgressMethods));
+            setInProgressMethods((oldInProgressMethods) => {
+              const methods =
+                InProgressMethods.fromExisting(oldInProgressMethods);
+              methods.setPackageMethods(
+                msg.packageName,
+                new Set(msg.inProgressMethods),
+              );
+              return methods;
+            });
             break;
           default:
             assertNever(msg);
@@ -346,7 +355,7 @@ export function DataExtensionsEditor({
           externalApiUsages={externalApiUsages}
           modeledMethods={modeledMethods}
           modifiedSignatures={modifiedSignatures}
-          inProgressSignatures={inProgressSignatures}
+          inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledApis={hideModeledApis}
           onChange={onChange}

extensions/ql-vscode/src/view/data-extensions-editor/LibraryRow.tsx

@@ -14,6 +14,7 @@ import {
   VSCodeTag,
 } from "@vscode/webview-ui-toolkit/react";
 import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 const LibraryContainer = styled.div`
   background-color: var(--vscode-peekViewResult-background);
@@ -72,7 +73,7 @@ type Props = {
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
-  inProgressSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   viewState: DataExtensionEditorViewState;
   hideModeledApis: boolean;
   onChange: (
@@ -100,7 +101,7 @@ export const LibraryRow = ({
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
-  inProgressSignatures,
+  inProgressMethods,
   viewState,
   hideModeledApis,
   onChange,
@@ -180,9 +181,9 @@ export const LibraryRow = ({
 
   const canStopAutoModeling = useMemo(() => {
     return externalApiUsages.some((externalApiUsage) =>
-      inProgressSignatures.has(externalApiUsage.signature),
+      inProgressMethods.hasMethod(title, externalApiUsage.signature),
     );
-  }, [externalApiUsages, inProgressSignatures]);
+  }, [externalApiUsages, title, inProgressMethods]);
 
   return (
     <LibraryContainer>
@@ -232,10 +233,11 @@ export const LibraryRow = ({
         <>
           <SectionDivider />
           <ModeledMethodDataGrid
+            packageName={title}
             externalApiUsages={externalApiUsages}
             modeledMethods={modeledMethods}
             modifiedSignatures={modifiedSignatures}
-            inProgressSignatures={inProgressSignatures}
+            inProgressMethods={inProgressMethods}
             mode={viewState.mode}
             hideModeledApis={hideModeledApis}
             onChange={onChangeWithModelName}

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodDataGrid.tsx

@@ -10,12 +10,14 @@ import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { useMemo } from "react";
 import { Mode } from "../../data-extensions-editor/shared/mode";
 import { sortMethods } from "../../data-extensions-editor/shared/sorting";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 type Props = {
+  packageName: string;
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
-  inProgressSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   mode: Mode;
   hideModeledApis: boolean;
   onChange: (
@@ -25,10 +27,11 @@ type Props = {
 };
 
 export const ModeledMethodDataGrid = ({
+  packageName,
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
-  inProgressSignatures,
+  inProgressMethods,
   mode,
   hideModeledApis,
   onChange,
@@ -63,7 +66,8 @@ export const ModeledMethodDataGrid = ({
           externalApiUsage={externalApiUsage}
           modeledMethod={modeledMethods[externalApiUsage.signature]}
           methodIsUnsaved={modifiedSignatures.has(externalApiUsage.signature)}
-          modelingInProgress={inProgressSignatures.has(
+          modelingInProgress={inProgressMethods.hasMethod(
+            packageName,
             externalApiUsage.signature,
           )}
           mode={mode}

extensions/ql-vscode/src/view/data-extensions-editor/ModeledMethodsList.tsx

@@ -9,12 +9,13 @@ import {
   sortGroupNames,
 } from "../../data-extensions-editor/shared/sorting";
 import { DataExtensionEditorViewState } from "../../data-extensions-editor/shared/view-state";
+import { InProgressMethods } from "../../data-extensions-editor/shared/in-progress-methods";
 
 type Props = {
   externalApiUsages: ExternalApiUsage[];
   modeledMethods: Record<string, ModeledMethod>;
   modifiedSignatures: Set<string>;
-  inProgressSignatures: Set<string>;
+  inProgressMethods: InProgressMethods;
   viewState: DataExtensionEditorViewState;
   hideModeledApis: boolean;
   onChange: (
@@ -44,7 +45,7 @@ export const ModeledMethodsList = ({
   externalApiUsages,
   modeledMethods,
   modifiedSignatures,
-  inProgressSignatures,
+  inProgressMethods,
   viewState,
   hideModeledApis,
   onChange,
@@ -89,7 +90,7 @@ export const ModeledMethodsList = ({
           externalApiUsages={grouped[libraryName]}
           modeledMethods={modeledMethods}
           modifiedSignatures={modifiedSignatures}
-          inProgressSignatures={inProgressSignatures}
+          inProgressMethods={inProgressMethods}
           viewState={viewState}
           hideModeledApis={hideModeledApis}
           onChange={onChange}