Remove method and modeledMethods from SaveModeledMethods message

Author: robertbrignull

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

@@ -543,8 +543,7 @@ interface RefreshMethods {
 
 interface SaveModeledMethods {
   t: "saveModeledMethods";
-  methods: Method[];
-  modeledMethods: Record<string, ModeledMethod>;
+  methodSignatures?: string[];
 }
 
 interface GenerateMethodMessage {

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

@@ -228,47 +228,58 @@ export class ModelEditorView extends AbstractWebview<
 
         break;
       case "saveModeledMethods":
-        await withProgress(
-          async (progress) => {
-            progress({
-              step: 1,
-              maxStep: 500 + externalApiQueriesProgressMaxStep,
-              message: "Writing model files",
-            });
-            await saveModeledMethods(
-              this.extensionPack,
-              this.databaseItem.language,
-              msg.methods,
-              msg.modeledMethods,
-              this.mode,
-              this.cliServer,
-              this.app.logger,
-            );
-
-            await Promise.all([
-              this.setViewState(),
-              this.loadMethods((update) =>
-                progress({
-                  ...update,
-                  step: update.step + 500,
-                  maxStep: 500 + externalApiQueriesProgressMaxStep,
-                }),
-              ),
-            ]);
-          },
-          {
-            cancellable: false,
-          },
-        );
+        {
+          const methods = this.modelingStore.getMethods(
+            this.databaseItem,
+            msg.methodSignatures,
+          );
+          const modeledMethods = this.modelingStore.getModeledMethods(
+            this.databaseItem,
+            msg.methodSignatures,
+          );
 
-        this.modelingStore.removeModifiedMethods(
-          this.databaseItem,
-          Object.keys(msg.modeledMethods),
-        );
+          await withProgress(
+            async (progress) => {
+              progress({
+                step: 1,
+                maxStep: 500 + externalApiQueriesProgressMaxStep,
+                message: "Writing model files",
+              });
+              await saveModeledMethods(
+                this.extensionPack,
+                this.databaseItem.language,
+                methods,
+                modeledMethods,
+                this.mode,
+                this.cliServer,
+                this.app.logger,
+              );
+
+              await Promise.all([
+                this.setViewState(),
+                this.loadMethods((update) =>
+                  progress({
+                    ...update,
+                    step: update.step + 500,
+                    maxStep: 500 + externalApiQueriesProgressMaxStep,
+                  }),
+                ),
+              ]);
+            },
+            {
+              cancellable: false,
+            },
+          );
 
-        void telemetryListener?.sendUIInteraction(
-          "model-editor-save-modeled-methods",
-        );
+          this.modelingStore.removeModifiedMethods(
+            this.databaseItem,
+            Object.keys(modeledMethods),
+          );
+
+          void telemetryListener?.sendUIInteraction(
+            "model-editor-save-modeled-methods",
+          );
+        }
 
         break;
       case "generateMethod":

extensions/ql-vscode/src/model-editor/modeling-store.ts

@@ -154,6 +154,23 @@ export class ModelingStore extends DisposableObject {
     return this.state.get(this.activeDb);
   }
 
+  /**
+   * Returns the methods for the given database item and method signatures.
+   * If no method signatures are provided, returns all methods.
+   */
+  public getMethods(
+    dbItem: DatabaseItem,
+    methodSignatures?: string[],
+  ): Method[] {
+    const methods = this.getState(dbItem).methods;
+    if (!methodSignatures) {
+      return methods;
+    }
+    return methods.filter((method) =>
+      methodSignatures.includes(method.signature),
+    );
+  }
+
   public setMethods(dbItem: DatabaseItem, methods: Method[]) {
     const dbState = this.getState(dbItem);
     const dbUri = dbItem.databaseUri.toString();
@@ -182,6 +199,25 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  /**
+   * Returns the modeled methods for the given database item and method signatures.
+   * If no method signatures are provided, returns all modeled methods.
+   */
+  public getModeledMethods(
+    dbItem: DatabaseItem,
+    methodSignatures?: string[],
+  ): Record<string, ModeledMethod> {
+    const modeledMethods = this.getState(dbItem).modeledMethods;
+    if (!methodSignatures) {
+      return modeledMethods;
+    }
+    return Object.fromEntries(
+      Object.entries(modeledMethods).filter(([key]) =>
+        methodSignatures.includes(key),
+      ),
+    );
+  }
+
   public addModeledMethods(
     dbItem: DatabaseItem,
     methods: Record<string, ModeledMethod>,

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

@@ -78,10 +78,7 @@ export type LibraryRowProps = {
   hideModeledMethods: boolean;
   revealedMethodSignature: string | null;
   onChange: (modeledMethod: ModeledMethod) => void;
-  onSaveModelClick: (
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
-  ) => void;
+  onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     dependencyName: string,
     methods: Method[],
@@ -165,11 +162,11 @@ export const LibraryRow = ({
 
   const handleSave = useCallback(
     async (e: React.MouseEvent) => {
-      onSaveModelClick(methods, modeledMethods);
+      onSaveModelClick(methods.map((m) => m.signature));
       e.stopPropagation();
       e.preventDefault();
     },
-    [methods, modeledMethods, onSaveModelClick],
+    [methods, onSaveModelClick],
   );
 
   const hasUnsavedChanges = useMemo(() => {

extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx

@@ -196,21 +196,15 @@ export function ModelEditor({
   const onSaveAllClick = useCallback(() => {
     vscode.postMessage({
       t: "saveModeledMethods",
-      methods,
-      modeledMethods,
     });
-  }, [methods, modeledMethods]);
+  }, []);
 
-  const onSaveModelClick = useCallback(
-    (methods: Method[], modeledMethods: Record<string, ModeledMethod>) => {
-      vscode.postMessage({
-        t: "saveModeledMethods",
-        methods,
-        modeledMethods,
-      });
-    },
-    [],
-  );
+  const onSaveModelClick = useCallback((methodSignatures: string[]) => {
+    vscode.postMessage({
+      t: "saveModeledMethods",
+      methodSignatures,
+    });
+  }, []);
 
   const onGenerateFromSourceClick = useCallback(() => {
     vscode.postMessage({

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

@@ -20,10 +20,7 @@ export type ModeledMethodsListProps = {
   viewState: ModelEditorViewState;
   hideModeledMethods: boolean;
   onChange: (modeledMethod: ModeledMethod) => void;
-  onSaveModelClick: (
-    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
-  ) => void;
+  onSaveModelClick: (methodSignatures: string[]) => void;
   onGenerateFromLlmClick: (
     packageName: string,
     methods: Method[],