Merge pull request #2906 from github/koesie10/convert-remaining-multiple-models

Convert remaining extension host code to handle multiple models
2023-10-10 11:43:09 +02:00
parent 4dcfa8b679 385b0e8d1a
commit 319a54c32f
11 changed files with 158 additions and 102 deletions
--- a/extensions/ql-vscode/src/model-editor/auto-model.ts
+++ b/extensions/ql-vscode/src/model-editor/auto-model.ts
@@ -14,13 +14,13 @@ import { groupMethods, sortGroupNames, sortMethods } from "./shared/sorting";
 * the order in the UI.
 * @param mode Whether it is application or framework mode.
 * @param methods all methods.
- * @param modeledMethods the currently modeled methods.
+ * @param modeledMethodsBySignature the currently modeled methods.
 * @returns list of modeled methods that are candidates for modeling.
 */
 export function getCandidates(
  mode: Mode,
  methods: Method[],
-  modeledMethods: Record<string, ModeledMethod>,
+  modeledMethodsBySignature: Record<string, ModeledMethod[]>,
 ): MethodSignature[] {
  // Sort the same way as the UI so we send the first ones listed in the UI first
  const grouped = groupMethods(methods, mode);
@@ -32,12 +32,11 @@ export function getCandidates(
  const candidates: MethodSignature[] = [];

  for (const method of sortedMethods) {
-    const modeledMethod: ModeledMethod = modeledMethods[method.signature] ?? {
-      type: "none",
-    };
+    const modeledMethods: ModeledMethod[] =
+      modeledMethodsBySignature[method.signature] ?? [];

    // Anything that is modeled is not a candidate
-    if (modeledMethod.type !== "none") {
+    if (modeledMethods.some((m) => m.type !== "none")) {
      continue;
    }

--- a/extensions/ql-vscode/src/model-editor/auto-modeler.ts
+++ b/extensions/ql-vscode/src/model-editor/auto-modeler.ts
@@ -16,7 +16,6 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
-import { convertToLegacyModeledMethods } from "./modeled-methods-legacy";

 // Limit the number of candidates we send to the model in each request
 // to avoid long requests.
@@ -43,7 +42,7 @@ export class AutoModeler {
      inProgressMethods: string[],
    ) => Promise<void>,
    private readonly addModeledMethods: (
-      modeledMethods: Record<string, ModeledMethod>,
+      modeledMethods: Record<string, ModeledMethod[]>,
    ) => Promise<void>,
  ) {
    this.jobs = new Map<string, CancellationTokenSource>();
@@ -60,7 +59,7 @@ export class AutoModeler {
  public async startModeling(
    packageName: string,
    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    mode: Mode,
  ): Promise<void> {
    if (this.jobs.has(packageName)) {
@@ -107,7 +106,7 @@ export class AutoModeler {
  private async modelPackage(
    packageName: string,
    methods: Method[],
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    mode: Mode,
    cancellationTokenSource: CancellationTokenSource,
  ): Promise<void> {
@@ -193,31 +192,31 @@ export class AutoModeler {
      filename: "auto-model.yml",
    });

-    const rawLoadedMethods = loadDataExtensionYaml(models);
-    if (!rawLoadedMethods) {
+    const loadedMethods = loadDataExtensionYaml(models);
+    if (!loadedMethods) {
      return;
    }

-    const loadedMethods = convertToLegacyModeledMethods(rawLoadedMethods);
-
    // Any candidate that was part of the response is a negative result
    // meaning that the canidate is not a sink for the kinds that the LLM is checking for.
    // For now we model this as a sink neutral method, however this is subject
    // to discussion.
    for (const candidate of candidateMethods) {
      if (!(candidate.signature in loadedMethods)) {
-        loadedMethods[candidate.signature] = {
-          type: "neutral",
-          kind: "sink",
-          input: "",
-          output: "",
-          provenance: "ai-generated",
-          signature: candidate.signature,
-          packageName: candidate.packageName,
-          typeName: candidate.typeName,
-          methodName: candidate.methodName,
-          methodParameters: candidate.methodParameters,
-        };
+        loadedMethods[candidate.signature] = [
+          {
+            type: "neutral",
+            kind: "sink",
+            input: "",
+            output: "",
+            provenance: "ai-generated",
+            signature: candidate.signature,
+            packageName: candidate.packageName,
+            typeName: candidate.typeName,
+            methodName: candidate.methodName,
+            methodParameters: candidate.methodParameters,
+          },
+        ];
      }
    }

--- a/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
+++ b/extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts
@@ -14,6 +14,10 @@ import { assertNever } from "../../common/helpers-pure";
 import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
 import { DatabaseItem } from "../../databases/local-databases";
+import {
+  convertFromLegacyModeledMethod,
+  convertToLegacyModeledMethod,
+} from "../modeled-methods-legacy";

 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
  ToMethodModelingMessage,
@@ -70,7 +74,9 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
        await this.postMessage({
          t: "setSelectedMethod",
          method: selectedMethod.method,
-          modeledMethod: selectedMethod.modeledMethod,
+          modeledMethod: convertToLegacyModeledMethod(
+            selectedMethod.modeledMethods,
+          ),
          isModified: selectedMethod.isModified,
        });
      }
@@ -107,9 +113,10 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
      case "setModeledMethod": {
        const activeState = this.ensureActiveState();

-        this.modelingStore.updateModeledMethod(
+        this.modelingStore.updateModeledMethods(
          activeState.databaseItem,
-          msg.method,
+          msg.method.signature,
+          convertFromLegacyModeledMethod(msg.method),
        );
        this.modelingStore.addModifiedMethod(
          activeState.databaseItem,
@@ -158,12 +165,15 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
    this.push(
      this.modelingStore.onModeledMethodsChanged(async (e) => {
        if (this.webviewView && e.isActiveDb) {
-          const modeledMethod = e.modeledMethods[this.method?.signature ?? ""];
-          if (modeledMethod) {
-            await this.postMessage({
-              t: "setModeledMethod",
-              method: modeledMethod,
-            });
+          const modeledMethods = e.modeledMethods[this.method?.signature ?? ""];
+          if (modeledMethods) {
+            const modeledMethod = convertToLegacyModeledMethod(modeledMethods);
+            if (modeledMethod) {
+              await this.postMessage({
+                t: "setModeledMethod",
+                method: modeledMethod,
+              });
+            }
          }
        }
      }),
@@ -190,7 +200,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
          await this.postMessage({
            t: "setSelectedMethod",
            method: e.method,
-            modeledMethod: e.modeledMethod,
+            modeledMethod: convertToLegacyModeledMethod(e.modeledMethods),
            isModified: e.isModified,
          });
        }
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
+++ b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts
@@ -26,7 +26,7 @@ export class MethodsUsageDataProvider
  private databaseItem: DatabaseItem | undefined = undefined;
  private sourceLocationPrefix: string | undefined = undefined;
  private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
-  private modeledMethods: Record<string, ModeledMethod> = {};
+  private modeledMethods: Record<string, ModeledMethod[]> = {};
  private modifiedMethodSignatures: Set<string> = new Set();

  private readonly onDidChangeTreeDataEmitter = this.push(
@@ -52,7 +52,7 @@ export class MethodsUsageDataProvider
    methods: Method[],
    databaseItem: DatabaseItem,
    hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    modifiedMethodSignatures: Set<string>,
  ): Promise<void> {
    if (
@@ -99,13 +99,10 @@ export class MethodsUsageDataProvider
  }

  private getModelingStatusIcon(method: Method): ThemeIcon {
-    const modeledMethod = this.modeledMethods[method.signature];
+    const modeledMethods = this.modeledMethods[method.signature];
    const modifiedMethod = this.modifiedMethodSignatures.has(method.signature);

-    const status = getModelingStatus(
-      modeledMethod ? [modeledMethod] : [],
-      modifiedMethod,
-    );
+    const status = getModelingStatus(modeledMethods, modifiedMethod);
    switch (status) {
      case "unmodeled":
        return new ThemeIcon("error", new ThemeColor("errorForeground"));
--- a/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
+++ b/extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts
@@ -34,7 +34,7 @@ export class MethodsUsagePanel extends DisposableObject {
    methods: Method[],
    databaseItem: DatabaseItem,
    hideModeledMethods: boolean,
-    modeledMethods: Record<string, ModeledMethod>,
+    modeledMethods: Record<string, ModeledMethod[]>,
    modifiedMethodSignatures: Set<string>,
  ): Promise<void> {
    await this.dataProvider.setState(
--- a/extensions/ql-vscode/src/model-editor/model-editor-view.ts
+++ b/extensions/ql-vscode/src/model-editor/model-editor-view.ts
@@ -44,7 +44,7 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
 import { ModelEditorViewTracker } from "./model-editor-view-tracker";
 import {
-  convertFromLegacyModeledMethods,
+  convertFromLegacyModeledMethod,
  convertToLegacyModeledMethods,
 } from "./modeled-methods-legacy";

@@ -224,7 +224,7 @@ export class ModelEditorView extends AbstractWebview<
                this.extensionPack,
                this.databaseItem.language,
                methods,
-                convertFromLegacyModeledMethods(modeledMethods),
+                modeledMethods,
                this.mode,
                this.cliServer,
                this.app.logger,
@@ -311,7 +311,10 @@ export class ModelEditorView extends AbstractWebview<
        );
        break;
      case "setModeledMethod": {
-        this.setModeledMethod(msg.method);
+        this.setModeledMethods(
+          msg.method.signature,
+          convertFromLegacyModeledMethod(msg.method),
+        );
        break;
      }
      default:
@@ -371,10 +374,7 @@ export class ModelEditorView extends AbstractWebview<
        this.cliServer,
        this.app.logger,
      );
-      this.modelingStore.setModeledMethods(
-        this.databaseItem,
-        convertToLegacyModeledMethods(modeledMethods),
-      );
+      this.modelingStore.setModeledMethods(this.databaseItem, modeledMethods);
    } catch (e: unknown) {
      void showAndLogErrorMessage(
        this.app.logger,
@@ -446,10 +446,16 @@ export class ModelEditorView extends AbstractWebview<
            queryStorageDir: this.queryStorageDir,
            databaseItem: addedDatabase ?? this.databaseItem,
            onResults: async (modeledMethods) => {
-              const modeledMethodsByName: Record<string, ModeledMethod> = {};
+              const modeledMethodsByName: Record<string, ModeledMethod[]> = {};

              for (const modeledMethod of modeledMethods) {
-                modeledMethodsByName[modeledMethod.signature] = modeledMethod;
+                if (!(modeledMethod.signature in modeledMethodsByName)) {
+                  modeledMethodsByName[modeledMethod.signature] = [];
+                }
+
+                modeledMethodsByName[modeledMethod.signature].push(
+                  modeledMethod,
+                );
              }

              this.addModeledMethods(modeledMethodsByName);
@@ -620,7 +626,7 @@ export class ModelEditorView extends AbstractWebview<
        if (event.dbUri === this.databaseItem.databaseUri.toString()) {
          await this.postMessage({
            t: "setModeledMethods",
-            methods: event.modeledMethods,
+            methods: convertToLegacyModeledMethods(event.modeledMethods),
          });
        }
      }),
@@ -646,7 +652,7 @@ export class ModelEditorView extends AbstractWebview<
    );
  }

-  private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
+  private addModeledMethods(modeledMethods: Record<string, ModeledMethod[]>) {
    this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);

    this.modelingStore.addModifiedMethods(
@@ -655,13 +661,17 @@ export class ModelEditorView extends AbstractWebview<
    );
  }

-  private setModeledMethod(method: ModeledMethod) {
+  private setModeledMethods(signature: string, methods: ModeledMethod[]) {
    const state = this.modelingStore.getStateForActiveDb();
    if (!state) {
      throw new Error("Attempting to set modeled method without active db");
    }

-    this.modelingStore.updateModeledMethod(state.databaseItem, method);
-    this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
+    this.modelingStore.updateModeledMethods(
+      state.databaseItem,
+      signature,
+      methods,
+    );
+    this.modelingStore.addModifiedMethod(state.databaseItem, signature);
  }
 }
--- a/extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts
+++ b/extensions/ql-vscode/src/model-editor/modeled-methods-legacy.ts
@@ -1,23 +1,57 @@
 import { ModeledMethod } from "./modeled-method";

-export function convertFromLegacyModeledMethods(
-  modeledMethods: Record<string, ModeledMethod>,
-): Record<string, ModeledMethod[]> {
-  // Convert a single ModeledMethod to an array of ModeledMethods
-  return Object.fromEntries(
-    Object.entries(modeledMethods).map(([signature, modeledMethod]) => {
-      return [signature, [modeledMethod]];
-    }),
-  );
-}
-
+/**
+ * Converts a record of a single ModeledMethod indexed by signature to a record of ModeledMethod[] indexed by signature
+ * for legacy usage. This function should always be used instead of the trivial conversion to track usages of this
+ * conversion.
+ *
+ * This method should only be called inside a `postMessage` call. If it's used anywhere else, consider whether the
+ * boundary is correct: the boundary should as close as possible to the extension host -> webview boundary.
+ *
+ * @param modeledMethods The record of a single ModeledMethod indexed by signature
+ */
 export function convertToLegacyModeledMethods(
  modeledMethods: Record<string, ModeledMethod[]>,
 ): Record<string, ModeledMethod> {
  // Always take the first modeled method in the array
  return Object.fromEntries(
-    Object.entries(modeledMethods).map(([signature, modeledMethods]) => {
-      return [signature, modeledMethods[0]];
-    }),
+    Object.entries(modeledMethods)
+      .map(([signature, modeledMethods]) => {
+        const modeledMethod = convertToLegacyModeledMethod(modeledMethods);
+        if (!modeledMethod) {
+          return null;
+        }
+        return [signature, modeledMethod];
+      })
+      .filter((entry): entry is [string, ModeledMethod] => entry !== null),
  );
 }
+
+/**
+ * Converts a single ModeledMethod to a ModeledMethod[] for legacy usage. This function should always be used instead
+ * of the trivial conversion to track usages of this conversion.
+ *
+ * This method should only be called inside a `onMessage` function (or its equivalent). If it's used anywhere else,
+ * consider whether the boundary is correct: the boundary should as close as possible to the webview -> extension host
+ * boundary.
+ *
+ * @param modeledMethod The single ModeledMethod
+ */
+export function convertFromLegacyModeledMethod(modeledMethod: ModeledMethod) {
+  return [modeledMethod];
+}
+
+/**
+ * Converts a ModeledMethod[] to a single ModeledMethod for legacy usage. This function should always be used instead
+ * of the trivial conversion to track usages of this conversion.
+ *
+ * This method should only be called inside a `postMessage` call. If it's used anywhere else, consider whether the
+ * boundary is correct: the boundary should as close as possible to the extension host -> webview boundary.
+ *
+ * @param modeledMethods The ModeledMethod[]
+ */
+export function convertToLegacyModeledMethod(
+  modeledMethods: ModeledMethod[],
+): ModeledMethod | undefined {
+  return modeledMethods[0];
+}
--- a/extensions/ql-vscode/src/model-editor/modeling-store.ts
+++ b/extensions/ql-vscode/src/model-editor/modeling-store.ts
@@ -10,7 +10,7 @@ export interface DbModelingState {
  databaseItem: DatabaseItem;
  methods: Method[];
  hideModeledMethods: boolean;
-  modeledMethods: Record<string, ModeledMethod>;
+  modeledMethods: Record<string, ModeledMethod[]>;
  modifiedMethodSignatures: Set<string>;
  selectedMethod: Method | undefined;
  selectedUsage: Usage | undefined;
@@ -28,7 +28,7 @@ interface HideModeledMethodsChangedEvent {
 }

 interface ModeledMethodsChangedEvent {
-  modeledMethods: Record<string, ModeledMethod>;
+  modeledMethods: Record<string, ModeledMethod[]>;
  dbUri: string;
  isActiveDb: boolean;
 }
@@ -43,7 +43,7 @@ interface SelectedMethodChangedEvent {
  databaseItem: DatabaseItem;
  method: Method;
  usage: Usage;
-  modeledMethod: ModeledMethod | undefined;
+  modeledMethods: ModeledMethod[];
  isModified: boolean;
 }

@@ -221,7 +221,7 @@ export class ModelingStore extends DisposableObject {
  public getModeledMethods(
    dbItem: DatabaseItem,
    methodSignatures?: string[],
-  ): Record<string, ModeledMethod> {
+  ): Record<string, ModeledMethod[]> {
    const modeledMethods = this.getState(dbItem).modeledMethods;
    if (!methodSignatures) {
      return modeledMethods;
@@ -235,14 +235,15 @@ export class ModelingStore extends DisposableObject {

  public addModeledMethods(
    dbItem: DatabaseItem,
-    methods: Record<string, ModeledMethod>,
+    methods: Record<string, ModeledMethod[]>,
  ) {
    this.changeModeledMethods(dbItem, (state) => {
      const newModeledMethods = {
        ...methods,
+        // Keep all methods that are already modeled in some form in the state
        ...Object.fromEntries(
-          Object.entries(state.modeledMethods).filter(
-            ([_, value]) => value.type !== "none",
+          Object.entries(state.modeledMethods).filter(([_, value]) =>
+            value.some((m) => m.type !== "none"),
          ),
        ),
      };
@@ -252,17 +253,21 @@ export class ModelingStore extends DisposableObject {

  public setModeledMethods(
    dbItem: DatabaseItem,
-    methods: Record<string, ModeledMethod>,
+    methods: Record<string, ModeledMethod[]>,
  ) {
    this.changeModeledMethods(dbItem, (state) => {
      state.modeledMethods = { ...methods };
    });
  }

-  public updateModeledMethod(dbItem: DatabaseItem, method: ModeledMethod) {
+  public updateModeledMethods(
+    dbItem: DatabaseItem,
+    signature: string,
+    modeledMethods: ModeledMethod[],
+  ) {
    this.changeModeledMethods(dbItem, (state) => {
      const newModeledMethods = { ...state.modeledMethods };
-      newModeledMethods[method.signature] = method;
+      newModeledMethods[signature] = modeledMethods;
      state.modeledMethods = newModeledMethods;
    });
  }
@@ -316,7 +321,7 @@ export class ModelingStore extends DisposableObject {
      databaseItem: dbItem,
      method,
      usage,
-      modeledMethod: dbState.modeledMethods[method.signature],
+      modeledMethods: dbState.modeledMethods[method.signature],
      isModified: dbState.modifiedMethodSignatures.has(method.signature),
    });
  }
@@ -335,7 +340,7 @@ export class ModelingStore extends DisposableObject {
    return {
      method: selectedMethod,
      usage: dbState.selectedUsage,
-      modeledMethod: dbState.modeledMethods[selectedMethod.signature],
+      modeledMethods: dbState.modeledMethods[selectedMethod.signature],
      isModified: dbState.modifiedMethodSignatures.has(
        selectedMethod.signature,
      ),
--- a/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts
+++ b/extensions/ql-vscode/test/unit-tests/model-editor/auto-model.test.ts
@@ -99,19 +99,21 @@ describe("getCandidates", () => {
        usages: [],
      },
    ];
-    const modeledMethods: Record<string, ModeledMethod> = {
-      "org.my.A#x()": {
-        type: "neutral",
-        kind: "",
-        input: "",
-        output: "",
-        provenance: "manual",
-        signature: "org.my.A#x()",
-        packageName: "org.my",
-        typeName: "A",
-        methodName: "x",
-        methodParameters: "()",
-      },
+    const modeledMethods: Record<string, ModeledMethod[]> = {
+      "org.my.A#x()": [
+        {
+          type: "neutral",
+          kind: "",
+          input: "",
+          output: "",
+          provenance: "manual",
+          signature: "org.my.A#x()",
+          packageName: "org.my",
+          typeName: "A",
+          methodName: "x",
+          methodParameters: "()",
+        },
+      ],
    };
    const candidates = getCandidates(Mode.Application, methods, modeledMethods);
    expect(candidates.length).toEqual(0);
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-data-provider.test.ts
@@ -20,7 +20,7 @@ describe("MethodsUsageDataProvider", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();
    const dbItem = mockedObject<DatabaseItem>({
      getSourceLocationPrefix: () => "test",
@@ -125,7 +125,7 @@ describe("MethodsUsageDataProvider", () => {
    });

    it("should emit onDidChangeTreeData event when modeled methods has changed", async () => {
-      const modeledMethods2: Record<string, ModeledMethod> = {};
+      const modeledMethods2: Record<string, ModeledMethod[]> = {};

      await dataProvider.setState(
        methods,
@@ -213,7 +213,7 @@ describe("MethodsUsageDataProvider", () => {
    });

    const methods: Method[] = [supportedMethod, unsupportedMethod];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();

    const dbItem = mockedObject<DatabaseItem>({
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/methods-usage/methods-usage-panel.test.ts
@@ -21,7 +21,7 @@ describe("MethodsUsagePanel", () => {
  describe("setState", () => {
    const hideModeledMethods = false;
    const methods: Method[] = [createMethod()];
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();

    it("should update the tree view with the correct batch number", async () => {
@@ -50,7 +50,7 @@ describe("MethodsUsagePanel", () => {
    let modelingStore: ModelingStore;

    const hideModeledMethods: boolean = false;
-    const modeledMethods: Record<string, ModeledMethod> = {};
+    const modeledMethods: Record<string, ModeledMethod[]> = {};
    const modifiedMethodSignatures: Set<string> = new Set();
    const usage = createUsage();