Merge branch 'koesie10/pick-extension-model-file' into koesie10/create-extension-model-file

extensions/ql-vscode/src/cli.ts

@@ -107,17 +107,18 @@ export type MlModelInfo = {
 /** The expected output of `codeql resolve ml-models`. */
 export type MlModelsInfo = { models: MlModelInfo[] };
 
-export type DataExtensionInfo = {
+/** 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 ExtensionsInfo = {
+export type ResolveExtensionsResult = {
   models: MlModelInfo[];
   data: {
-    [filename: string]: DataExtensionInfo[];
+    [path: string]: DataExtensionResult[];
   };
 };
 
@@ -1215,11 +1216,11 @@ export class CodeQLCliServer implements Disposable {
   async resolveExtensions(
     suite: string,
     additionalPacks: string[],
-  ): Promise<ExtensionsInfo> {
+  ): Promise<ResolveExtensionsResult> {
     const args = this.getAdditionalPacksArg(additionalPacks);
     args.push(suite);
 
-    return this.runJsonCodeQlCliCommand<ExtensionsInfo>(
+    return this.runJsonCodeQlCliCommand<ResolveExtensionsResult>(
       ["resolve", "extensions"],
       args,
       "Resolving extensions",

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

@@ -9,7 +9,7 @@ import { join } from "path";
 import { App } from "../common/app";
 import { showAndLogErrorMessage } from "../helpers";
 import { withProgress } from "../progress";
-import { pickExtensionPack, pickModelFile } from "./extension-packs";
+import { pickExtensionPackModelFile } from "./extension-pack-picker";
 
 export class DataExtensionsEditorModule {
   private readonly queryStorageDir: string;
@@ -67,20 +67,11 @@ export class DataExtensionsEditorModule {
               return;
             }
 
-            const extensionPackPath = await pickExtensionPack(
-              this.cliServer,
-              progress,
-            );
-            if (!extensionPackPath) {
-              return;
-            }
-
-            const modelFile = await pickModelFile(
+            const modelFile = await pickExtensionPackModelFile(
               this.cliServer,
+              db,
               progress,
               token,
-              db,
-              extensionPackPath,
             );
             if (!modelFile) {
               return;

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

@@ -1,4 +1,4 @@
-import { relative, resolve } from "path";
+import { relative, resolve, sep } from "path";
 import { pathExists, readFile } from "fs-extra";
 import { load as loadYaml } from "js-yaml";
 import { minimatch } from "minimatch";
@@ -9,14 +9,42 @@ import { ProgressCallback } from "../progress";
 import { DatabaseItem } from "../local-databases";
 import { getQlPackPath, QLPACK_FILENAMES } from "../pure/ql";
 
-export async function pickExtensionPack(
-  cliServer: CodeQLCliServer,
+const maxStep = 3;
+
+export async function pickExtensionPackModelFile(
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks" | "resolveExtensions">,
+  databaseItem: Pick<DatabaseItem, "name">,
   progress: ProgressCallback,
+  token: CancellationToken,
+): Promise<string | undefined> {
+  const extensionPackPath = await pickExtensionPack(cliServer, progress, token);
+  if (!extensionPackPath) {
+    return;
+  }
+
+  const modelFile = await pickModelFile(
+    cliServer,
+    databaseItem,
+    extensionPackPath,
+    progress,
+    token,
+  );
+  if (!modelFile) {
+    return;
+  }
+
+  return modelFile;
+}
+
+async function pickExtensionPack(
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks">,
+  progress: ProgressCallback,
+  token: CancellationToken,
 ): Promise<string | undefined> {
   progress({
     message: "Resolving extension packs...",
     step: 1,
-    maxStep: 3,
+    maxStep,
   });
 
   // Get all existing extension packs in the workspace
@@ -30,12 +58,16 @@ export async function pickExtensionPack(
   progress({
     message: "Choosing extension pack...",
     step: 2,
-    maxStep: 3,
+    maxStep,
   });
 
-  const extensionPackOption = await window.showQuickPick(options, {
-    title: "Select extension pack to use",
-  });
+  const extensionPackOption = await window.showQuickPick(
+    options,
+    {
+      title: "Select extension pack to use",
+    },
+    token,
+  );
   if (!extensionPackOption) {
     return undefined;
   }
@@ -44,6 +76,13 @@ export async function pickExtensionPack(
   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;
   }
@@ -51,12 +90,12 @@ export async function pickExtensionPack(
   return extensionPackPaths[0];
 }
 
-export async function pickModelFile(
-  cliServer: CodeQLCliServer,
+async function pickModelFile(
+  cliServer: Pick<CodeQLCliServer, "resolveExtensions">,
+  databaseItem: Pick<DatabaseItem, "name">,
+  extensionPackPath: string,
   progress: ProgressCallback,
   token: CancellationToken,
-  databaseItem: DatabaseItem,
-  extensionPackPath: string,
 ): Promise<string | undefined> {
   // Find the existing model files in the extension pack
   const additionalPacks = getOnDiskWorkspaceFolders();
@@ -74,13 +113,13 @@ export async function pickModelFile(
   }
 
   if (modelFiles.size === 0) {
-    return pickNewModelFile(token, databaseItem, extensionPackPath);
+    return pickNewModelFile(databaseItem, extensionPackPath, token);
   }
 
   const fileOptions: Array<{ label: string; file: string | null }> = [];
   for (const file of modelFiles) {
     fileOptions.push({
-      label: relative(extensionPackPath, file),
+      label: relative(extensionPackPath, file).replaceAll(sep, "/"),
       file,
     });
   }
@@ -92,7 +131,7 @@ export async function pickModelFile(
   progress({
     message: "Choosing model file...",
     step: 3,
-    maxStep: 3,
+    maxStep,
   });
 
   const fileOption = await window.showQuickPick(
@@ -111,13 +150,13 @@ export async function pickModelFile(
     return fileOption.file;
   }
 
-  return pickNewModelFile(token, databaseItem, extensionPackPath);
+  return pickNewModelFile(databaseItem, extensionPackPath, token);
 }
 
 async function pickNewModelFile(
-  token: CancellationToken,
-  databaseItem: DatabaseItem,
+  databaseItem: Pick<DatabaseItem, "name">,
   extensionPackPath: string,
+  token: CancellationToken,
 ) {
   const qlpackPath = await getQlPackPath(extensionPackPath);
   if (!qlpackPath) {

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

@@ -0,0 +1,243 @@
+import { CancellationTokenSource, 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 databaseItem = {
+    name: "github/vscode-codeql",
+  };
+
+  const cancellationTokenSource = new CancellationTokenSource();
+  const token = cancellationTokenSource.token;
+
+  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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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,
+        databaseItem,
+        progress,
+        token,
+      ),
+    ).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),
+  };
+}