Add configurable model filename to data extension editor

This adds a pickable model filename from an existing extension pack to
the data extensions editor. This allows the user to edit one of their
existing data extensions. This does not yet add the ability to create
new extension packs and/or new model files.

This uses the `codeql resolve extensions` command to get the list of
available model files. This should be available in all CLI versions
which the data extensions editor supports.

Author: koesie10

extensions/ql-vscode/src/cli.ts

@@ -107,6 +107,20 @@ export type MlModelInfo = {
 /** The expected output of `codeql resolve ml-models`. */
 export type MlModelsInfo = { models: MlModelInfo[] };
 
+export type DataExtensionInfo = {
+  predicate: string;
+  file: string;
+  index: number;
+};
+
+/** The expected output of `codeql resolve extensions`. */
+export type ExtensionsInfo = {
+  models: MlModelInfo[];
+  data: {
+    [filename: string]: DataExtensionInfo[];
+  };
+};
+
 /**
  * The expected output of `codeql resolve qlref`.
  */
@@ -1192,6 +1206,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<ExtensionsInfo> {
+    const args = this.getAdditionalPacksArg(additionalPacks);
+    args.push(suite);
+
+    return this.runJsonCodeQlCliCommand<ExtensionsInfo>(
+      ["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 { pickExtensionPack, pickModelFile } from "./extension-packs";
 
 export class DataExtensionsEditorModule {
   private readonly queryStorageDir: string;
@@ -49,31 +51,55 @@ 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 extensionPackPath = await pickExtensionPack(
+              this.cliServer,
+              progress,
+            );
+            if (!extensionPackPath) {
+              return;
+            }
+
+            const modelFile = await pickModelFile(
+              this.cliServer,
+              progress,
+              extensionPackPath,
+            );
+            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-packs.ts

@@ -0,0 +1,92 @@
+import { relative } from "path";
+import { window } from "vscode";
+import { CodeQLCliServer } from "../cli";
+import { getOnDiskWorkspaceFolders, showAndLogErrorMessage } from "../helpers";
+import { ProgressCallback } from "../progress";
+
+export async function pickExtensionPack(
+  cliServer: CodeQLCliServer,
+  progress: ProgressCallback,
+): Promise<string | undefined> {
+  progress({
+    message: "Resolving extension packs...",
+    step: 1,
+    maxStep: 3,
+  });
+
+  // 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: 3,
+  });
+
+  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`,
+    );
+    return undefined;
+  }
+
+  return extensionPackPaths[0];
+}
+
+export async function pickModelFile(
+  cliServer: CodeQLCliServer,
+  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),
+      file,
+    });
+  }
+
+  progress({
+    message: "Choosing model file...",
+    step: 3,
+    maxStep: 3,
+  });
+
+  const fileOption = await window.showQuickPick(fileOptions, {
+    title: "Select model file to use",
+  });
+
+  if (!fileOption) {
+    return;
+  }
+
+  return fileOption.file;
+}