Merge pull request #2295 from github/koesie10/pick-extension-model-file

Add configurable model filename to data extension editor

extensions/ql-vscode/src/cli.ts

@@ -107,6 +107,21 @@ export type MlModelInfo = {
 /** The expected output of `codeql resolve ml-models`. */
 export type MlModelsInfo = { models: MlModelInfo[] };
 
+/** Information about a data extension predicate, as resolved by `codeql resolve extensions`. */
+export type DataExtensionResult = {
+  predicate: string;
+  file: string;
+  index: number;
+};
+
+/** The expected output of `codeql resolve extensions`. */
+export type ResolveExtensionsResult = {
+  models: MlModelInfo[];
+  data: {
+    [path: string]: DataExtensionResult[];
+  };
+};
+
 /**
  * The expected output of `codeql resolve qlref`.
  */
@@ -1192,6 +1207,29 @@ export class CodeQLCliServer implements Disposable {
     );
   }
 
+  /**
+   * Gets information about available extensions
+   * @param suite The suite to resolve.
+   * @param additionalPacks A list of directories to search for qlpacks.
+   * @returns An object containing the list of models and extensions
+   */
+  async resolveExtensions(
+    suite: string,
+    additionalPacks: string[],
+  ): Promise<ResolveExtensionsResult> {
+    const args = this.getAdditionalPacksArg(additionalPacks);
+    args.push(suite);
+
+    return this.runJsonCodeQlCliCommand<ResolveExtensionsResult>(
+      ["resolve", "extensions"],
+      args,
+      "Resolving extensions",
+      {
+        addFormat: false,
+      },
+    );
+  }
+
   /**
    * Gets information about the available languages.
    * @returns A dictionary mapping language name to the directory it comes from

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

@@ -8,6 +8,8 @@ import { ensureDir } from "fs-extra";
 import { join } from "path";
 import { App } from "../common/app";
 import { showAndLogErrorMessage } from "../helpers";
+import { withProgress } from "../progress";
+import { pickExtensionPackModelFile } from "./extension-pack-picker";
 
 export class DataExtensionsEditorModule {
   private readonly queryStorageDir: string;
@@ -49,31 +51,46 @@ export class DataExtensionsEditorModule {
 
   public getCommands(): DataExtensionsEditorCommands {
     return {
-      "codeQL.openDataExtensionsEditor": async () => {
-        const db = this.databaseManager.currentDatabaseItem;
-        if (!db) {
-          void showAndLogErrorMessage("No database selected");
-          return;
-        }
+      "codeQL.openDataExtensionsEditor": async () =>
+        withProgress(
+          async (progress) => {
+            const db = this.databaseManager.currentDatabaseItem;
+            if (!db) {
+              void showAndLogErrorMessage("No database selected");
+              return;
+            }
 
-        if (!(await this.cliServer.cliConstraints.supportsQlpacksKind())) {
-          void showAndLogErrorMessage(
-            `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_QLPACKS_KIND.format()} or later.`,
-          );
-          return;
-        }
+            if (!(await this.cliServer.cliConstraints.supportsQlpacksKind())) {
+              void showAndLogErrorMessage(
+                `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_QLPACKS_KIND.format()} or later.`,
+              );
+              return;
+            }
 
-        const view = new DataExtensionsEditorView(
-          this.ctx,
-          this.app,
-          this.databaseManager,
-          this.cliServer,
-          this.queryRunner,
-          this.queryStorageDir,
-          db,
-        );
-        await view.openView();
-      },
+            const modelFile = await pickExtensionPackModelFile(
+              this.cliServer,
+              progress,
+            );
+            if (!modelFile) {
+              return;
+            }
+
+            const view = new DataExtensionsEditorView(
+              this.ctx,
+              this.app,
+              this.databaseManager,
+              this.cliServer,
+              this.queryRunner,
+              this.queryStorageDir,
+              db,
+              modelFile,
+            );
+            await view.openView();
+          },
+          {
+            title: "Opening Data Extensions Editor",
+          },
+        ),
     };
   }
 

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

@@ -1,7 +1,6 @@
 import {
   CancellationTokenSource,
   ExtensionContext,
-  Uri,
   ViewColumn,
   window,
   workspace,
@@ -61,6 +60,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
     private readonly queryRunner: QueryRunner,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
+    private readonly modelFilename: string,
   ) {
     super(ctx);
   }
@@ -148,29 +148,19 @@ export class DataExtensionsEditorView extends AbstractWebview<
     externalApiUsages: ExternalApiUsage[],
     modeledMethods: Record<string, ModeledMethod>,
   ): Promise<void> {
-    const modelFilename = this.calculateModelFilename();
-    if (!modelFilename) {
-      return;
-    }
-
     const yaml = createDataExtensionYaml(externalApiUsages, modeledMethods);
 
-    await writeFile(modelFilename, yaml);
+    await writeFile(this.modelFilename, yaml);
 
-    void extLogger.log(`Saved data extension YAML to ${modelFilename}`);
+    void extLogger.log(`Saved data extension YAML to ${this.modelFilename}`);
   }
 
   protected async loadExistingModeledMethods(): Promise<void> {
-    const modelFilename = this.calculateModelFilename();
-    if (!modelFilename) {
-      return;
-    }
-
     try {
-      const yaml = await readFile(modelFilename, "utf8");
+      const yaml = await readFile(this.modelFilename, "utf8");
 
       const data = loadYaml(yaml, {
-        filename: modelFilename,
+        filename: this.modelFilename,
       });
 
       const existingModeledMethods = loadDataExtensionYaml(data);
@@ -365,17 +355,4 @@ export class DataExtensionsEditorView extends AbstractWebview<
       message: "",
     });
   }
-
-  private calculateModelFilename(): string | undefined {
-    const workspaceFolder = getQlSubmoduleFolder();
-    if (!workspaceFolder) {
-      return;
-    }
-
-    return Uri.joinPath(
-      workspaceFolder.uri,
-      "java/ql/lib/ext",
-      `${this.databaseItem.name.replaceAll("/", ".")}.model.yml`,
-    ).fsPath;
-  }
 }

extensions/ql-vscode/src/data-extensions-editor/extension-pack-picker.ts

@@ -0,0 +1,118 @@
+import { relative, sep } from "path";
+import { window } from "vscode";
+import { CodeQLCliServer } from "../cli";
+import { getOnDiskWorkspaceFolders, showAndLogErrorMessage } from "../helpers";
+import { ProgressCallback } from "../progress";
+
+const maxStep = 3;
+
+export async function pickExtensionPackModelFile(
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks" | "resolveExtensions">,
+  progress: ProgressCallback,
+): Promise<string | undefined> {
+  const extensionPackPath = await pickExtensionPack(cliServer, progress);
+  if (!extensionPackPath) {
+    return;
+  }
+
+  const modelFile = await pickModelFile(cliServer, progress, extensionPackPath);
+  if (!modelFile) {
+    return;
+  }
+
+  return modelFile;
+}
+
+async function pickExtensionPack(
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks">,
+  progress: ProgressCallback,
+): Promise<string | undefined> {
+  progress({
+    message: "Resolving extension packs...",
+    step: 1,
+    maxStep,
+  });
+
+  // Get all existing extension packs in the workspace
+  const additionalPacks = getOnDiskWorkspaceFolders();
+  const extensionPacks = await cliServer.resolveQlpacks(additionalPacks, true);
+  const options = Object.keys(extensionPacks).map((pack) => ({
+    label: pack,
+    extensionPack: pack,
+  }));
+
+  progress({
+    message: "Choosing extension pack...",
+    step: 2,
+    maxStep,
+  });
+
+  const extensionPackOption = await window.showQuickPick(options, {
+    title: "Select extension pack to use",
+  });
+  if (!extensionPackOption) {
+    return undefined;
+  }
+
+  const extensionPackPaths = extensionPacks[extensionPackOption.extensionPack];
+  if (extensionPackPaths.length !== 1) {
+    void showAndLogErrorMessage(
+      `Extension pack ${extensionPackOption.extensionPack} could not be resolved to a single location`,
+      {
+        fullMessage: `Extension pack ${
+          extensionPackOption.extensionPack
+        } could not be resolved to a single location. Found ${
+          extensionPackPaths.length
+        } locations: ${extensionPackPaths.join(", ")}.`,
+      },
+    );
+    return undefined;
+  }
+
+  return extensionPackPaths[0];
+}
+
+async function pickModelFile(
+  cliServer: Pick<CodeQLCliServer, "resolveExtensions">,
+  progress: ProgressCallback,
+  extensionPackPath: string,
+): Promise<string | undefined> {
+  // Find the existing model files in the extension pack
+  const additionalPacks = getOnDiskWorkspaceFolders();
+  const extensions = await cliServer.resolveExtensions(
+    extensionPackPath,
+    additionalPacks,
+  );
+
+  const modelFiles = new Set<string>();
+
+  if (extensionPackPath in extensions.data) {
+    for (const extension of extensions.data[extensionPackPath]) {
+      modelFiles.add(extension.file);
+    }
+  }
+
+  const fileOptions: Array<{ label: string; file: string }> = [];
+  for (const file of modelFiles) {
+    fileOptions.push({
+      label: relative(extensionPackPath, file).replaceAll(sep, "/"),
+      file,
+    });
+  }
+
+  progress({
+    message: "Choosing model file...",
+    step: 3,
+    maxStep,
+  });
+
+  const fileOption = await window.showQuickPick(fileOptions, {
+    title: "Select model file to use",
+  });
+
+  if (!fileOption) {
+    return;
+  }
+
+  return fileOption.file;
+}

extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/extension-pack-picker.test.ts

@@ -0,0 +1,207 @@
+import { QuickPickItem, window } from "vscode";
+
+import { pickExtensionPackModelFile } from "../../../../src/data-extensions-editor/extension-pack-picker";
+import { QlpacksInfo, ResolveExtensionsResult } from "../../../../src/cli";
+import * as helpers from "../../../../src/helpers";
+
+describe("pickExtensionPackModelFile", () => {
+  const qlPacks = {
+    "my-extension-pack": ["/a/b/c/my-extension-pack"],
+    "another-extension-pack": ["/a/b/c/another-extension-pack"],
+  };
+  const extensions = {
+    models: [],
+    data: {
+      "/a/b/c/my-extension-pack": [
+        {
+          file: "/a/b/c/my-extension-pack/models/model.yml",
+          index: 0,
+          predicate: "sinkModel",
+        },
+      ],
+    },
+  };
+
+  const progress = jest.fn();
+  let showQuickPickSpy: jest.SpiedFunction<typeof window.showQuickPick>;
+
+  beforeEach(() => {
+    showQuickPickSpy = jest
+      .spyOn(window, "showQuickPick")
+      .mockRejectedValue(new Error("Unexpected call to showQuickPick"));
+  });
+
+  it("allows choosing an existing extension pack and model file", async () => {
+    const cliServer = mockCliServer(qlPacks, extensions);
+
+    showQuickPickSpy.mockResolvedValueOnce({
+      label: "my-extension-pack",
+      extensionPack: "my-extension-pack",
+    } as QuickPickItem);
+    showQuickPickSpy.mockResolvedValueOnce({
+      label: "models/model.yml",
+      file: "/a/b/c/my-extension-pack/models/model.yml",
+    } as QuickPickItem);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      "/a/b/c/my-extension-pack/models/model.yml",
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledTimes(2);
+    expect(showQuickPickSpy).toHaveBeenCalledWith(
+      [
+        {
+          label: "my-extension-pack",
+          extensionPack: "my-extension-pack",
+        },
+        {
+          label: "another-extension-pack",
+          extensionPack: "another-extension-pack",
+        },
+      ],
+      {
+        title: expect.any(String),
+      },
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledWith(
+      [
+        {
+          label: "models/model.yml",
+          file: "/a/b/c/my-extension-pack/models/model.yml",
+        },
+      ],
+      {
+        title: expect.any(String),
+      },
+    );
+    expect(cliServer.resolveQlpacks).toHaveBeenCalledTimes(1);
+    expect(cliServer.resolveQlpacks).toHaveBeenCalledWith([], true);
+    expect(cliServer.resolveExtensions).toHaveBeenCalledTimes(1);
+    expect(cliServer.resolveExtensions).toHaveBeenCalledWith(
+      "/a/b/c/my-extension-pack",
+      [],
+    );
+  });
+
+  it("allows cancelling the extension pack prompt", async () => {
+    const cliServer = mockCliServer(qlPacks, extensions);
+
+    showQuickPickSpy.mockResolvedValueOnce(undefined);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      undefined,
+    );
+    expect(cliServer.resolveQlpacks).toHaveBeenCalled();
+    expect(cliServer.resolveExtensions).not.toHaveBeenCalled();
+  });
+
+  it("does not show any options when there are no extension packs", async () => {
+    const cliServer = mockCliServer({}, { models: [], data: {} });
+
+    showQuickPickSpy.mockResolvedValueOnce(undefined);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      undefined,
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledTimes(1);
+    expect(showQuickPickSpy).toHaveBeenCalledWith([], {
+      title: expect.any(String),
+    });
+    expect(cliServer.resolveQlpacks).toHaveBeenCalled();
+    expect(cliServer.resolveExtensions).not.toHaveBeenCalled();
+  });
+
+  it("shows an error when an extension pack resolves to more than 1 location", async () => {
+    const showAndLogErrorMessageSpy = jest.spyOn(
+      helpers,
+      "showAndLogErrorMessage",
+    );
+
+    const cliServer = mockCliServer(
+      {
+        "my-extension-pack": [
+          "/a/b/c/my-extension-pack",
+          "/a/b/c/my-extension-pack2",
+        ],
+      },
+      { models: [], data: {} },
+    );
+
+    showQuickPickSpy.mockResolvedValueOnce({
+      label: "my-extension-pack",
+      extensionPack: "my-extension-pack",
+    } as QuickPickItem);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      undefined,
+    );
+    expect(showAndLogErrorMessageSpy).toHaveBeenCalledTimes(1);
+    expect(showAndLogErrorMessageSpy).toHaveBeenCalledWith(
+      expect.stringMatching(/could not be resolved to a single location/),
+      expect.anything(),
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledTimes(1);
+    expect(cliServer.resolveQlpacks).toHaveBeenCalled();
+    expect(cliServer.resolveExtensions).not.toHaveBeenCalled();
+  });
+
+  it("allows cancelling the model file prompt", async () => {
+    const cliServer = mockCliServer(qlPacks, extensions);
+
+    showQuickPickSpy.mockResolvedValueOnce({
+      label: "my-extension-pack",
+      extensionPack: "my-extension-pack",
+    } as QuickPickItem);
+    showQuickPickSpy.mockResolvedValueOnce(undefined);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      undefined,
+    );
+    expect(cliServer.resolveQlpacks).toHaveBeenCalled();
+    expect(cliServer.resolveExtensions).toHaveBeenCalled();
+  });
+
+  it("does not show any options when there are no model files", async () => {
+    const cliServer = mockCliServer(qlPacks, { models: [], data: {} });
+
+    showQuickPickSpy.mockResolvedValueOnce({
+      label: "my-extension-pack",
+      extensionPack: "my-extension-pack",
+    } as QuickPickItem);
+    showQuickPickSpy.mockResolvedValueOnce(undefined);
+
+    expect(await pickExtensionPackModelFile(cliServer, progress)).toEqual(
+      undefined,
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledTimes(2);
+    expect(showQuickPickSpy).toHaveBeenCalledWith(
+      [
+        {
+          label: "my-extension-pack",
+          extensionPack: "my-extension-pack",
+        },
+        {
+          label: "another-extension-pack",
+          extensionPack: "another-extension-pack",
+        },
+      ],
+      {
+        title: expect.any(String),
+      },
+    );
+    expect(showQuickPickSpy).toHaveBeenCalledWith([], {
+      title: expect.any(String),
+    });
+    expect(cliServer.resolveQlpacks).toHaveBeenCalled();
+    expect(cliServer.resolveExtensions).toHaveBeenCalled();
+  });
+});
+
+function mockCliServer(
+  qlpacks: QlpacksInfo,
+  extensions: ResolveExtensionsResult,
+) {
+  return {
+    resolveQlpacks: jest.fn().mockResolvedValue(qlpacks),
+    resolveExtensions: jest.fn().mockResolvedValue(extensions),
+  };
+}