Merge pull request #2605 from github/robertbrignull/data-modeled-methods-tests

robertbrignull · web-flow · commit 5c12a4b205a6 · 2023-07-18T14:34:21.000+01:00
Refactor the code for loading/saving modeled methods to disk, and add tests
diff --git a/extensions/ql-vscode/src/codeql-cli/cli.ts b/extensions/ql-vscode/src/codeql-cli/cli.ts
@@ -1806,6 +1806,11 @@ export class CliVersionConstraint {
     "2.10.0",
   );
 
+  /**
+   * CLI version where the `resolve extensions` subcommand exists.
+   */
+  public static CLI_VERSION_WITH_RESOLVE_EXTENSIONS = new SemVer("2.10.2");
+
   /**
    * CLI version where the `--evaluator-log` and related options to the query server were introduced,
    * on a per-query server basis.
@@ -1882,6 +1887,12 @@ export class CliVersionConstraint {
     );
   }
 
+  async supportsResolveExtensions() {
+    return this.isVersionAtLeast(
+      CliVersionConstraint.CLI_VERSION_WITH_RESOLVE_EXTENSIONS,
+    );
+  }
+
   async supportsStructuredEvalLog() {
     return this.isVersionAtLeast(
       CliVersionConstraint.CLI_VERSION_WITH_STRUCTURED_EVAL_LOG,
diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-module.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-module.ts
@@ -78,6 +78,16 @@ export class DataExtensionsEditorModule {
               return;
             }
 
+            if (
+              !(await this.cliServer.cliConstraints.supportsResolveExtensions())
+            ) {
+              void showAndLogErrorMessage(
+                this.app.logger,
+                `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_RESOLVE_EXTENSIONS.format()} or later.`,
+              );
+              return;
+            }
+
             const modelFile = await pickExtensionPack(
               this.cliServer,
               db,
diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts
@@ -5,7 +5,6 @@ import {
   ViewColumn,
   window,
 } from "vscode";
-import { join } from "path";
 import { RequestError } from "@octokit/request-error";
 import {
   AbstractWebview,
@@ -21,8 +20,6 @@ import {
   showAndLogExceptionWithTelemetry,
   showAndLogErrorMessage,
 } from "../common/logging";
-import { outputFile, readFile } from "fs-extra";
-import { load as loadYaml } from "js-yaml";
 import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { asError, assertNever, getErrorMessage } from "../common/helpers-pure";
@@ -34,11 +31,6 @@ import { showResolvableLocation } from "../databases/local-databases/locations";
 import { decodeBqrsToExternalApiUsages } from "./bqrs";
 import { redactableError } from "../common/errors";
 import { readQueryResults, runQuery } from "./external-api-usage-query";
-import {
-  createDataExtensionYamlsForApplicationMode,
-  createDataExtensionYamlsForFrameworkMode,
-  loadDataExtensionYaml,
-} from "./yaml";
 import { ExternalApiUsage } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
@@ -49,8 +41,9 @@ import {
 } from "./auto-model";
 import { enableFrameworkMode, showLlmGeneration } from "../config";
 import { getAutoModelUsages } from "./auto-model-usages-query";
-import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { Mode } from "./shared/mode";
+import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
+import { join } from "path";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
@@ -123,9 +116,14 @@ export class DataExtensionsEditorView extends AbstractWebview<
 
         break;
       case "saveModeledMethods":
-        await this.saveModeledMethods(
+        await saveModeledMethods(
+          this.extensionPack,
+          this.databaseItem.name,
+          this.databaseItem.language,
           msg.externalApiUsages,
           msg.modeledMethods,
+          this.mode,
+          this.app.logger,
         );
         await Promise.all([this.setViewState(), this.loadExternalApiUsages()]);
 
@@ -194,79 +192,16 @@ export class DataExtensionsEditorView extends AbstractWebview<
     }
   }
 
-  protected async saveModeledMethods(
-    externalApiUsages: ExternalApiUsage[],
-    modeledMethods: Record<string, ModeledMethod>,
-  ): Promise<void> {
-    let yamls: Record<string, string>;
-    switch (this.mode) {
-      case Mode.Application:
-        yamls = createDataExtensionYamlsForApplicationMode(
-          this.databaseItem.language,
-          externalApiUsages,
-          modeledMethods,
-        );
-        break;
-      case Mode.Framework:
-        yamls = createDataExtensionYamlsForFrameworkMode(
-          this.databaseItem.name,
-          this.databaseItem.language,
-          externalApiUsages,
-          modeledMethods,
-        );
-        break;
-      default:
-        assertNever(this.mode);
-    }
-
-    for (const [filename, yaml] of Object.entries(yamls)) {
-      await outputFile(join(this.extensionPack.path, filename), yaml);
-    }
-
-    void this.app.logger.log(`Saved data extension YAML`);
-  }
-
   protected async loadExistingModeledMethods(): Promise<void> {
     try {
-      const extensions = await this.cliServer.resolveExtensions(
-        this.extensionPack.path,
-        getOnDiskWorkspaceFolders(),
+      const modeledMethods = await loadModeledMethods(
+        this.extensionPack,
+        this.cliServer,
+        this.app.logger,
       );
-
-      const modelFiles = new Set<string>();
-
-      if (this.extensionPack.path in extensions.data) {
-        for (const extension of extensions.data[this.extensionPack.path]) {
-          modelFiles.add(extension.file);
-        }
-      }
-
-      const existingModeledMethods: Record<string, ModeledMethod> = {};
-
-      for (const modelFile of modelFiles) {
-        const yaml = await readFile(modelFile, "utf8");
-
-        const data = loadYaml(yaml, {
-          filename: modelFile,
-        });
-
-        const modeledMethods = loadDataExtensionYaml(data);
-        if (!modeledMethods) {
-          void showAndLogErrorMessage(
-            this.app.logger,
-            `Failed to parse data extension YAML ${modelFile}.`,
-          );
-          continue;
-        }
-
-        for (const [key, value] of Object.entries(modeledMethods)) {
-          existingModeledMethods[key] = value;
-        }
-      }
-
       await this.postMessage({
         t: "loadModeledMethods",
-        modeledMethods: existingModeledMethods,
+        modeledMethods,
       });
     } catch (e: unknown) {
       void showAndLogErrorMessage(
diff --git a/extensions/ql-vscode/src/data-extensions-editor/modeled-method-fs.ts b/extensions/ql-vscode/src/data-extensions-editor/modeled-method-fs.ts
@@ -0,0 +1,93 @@
+import { outputFile, readFile } from "fs-extra";
+import { ExternalApiUsage } from "./external-api-usage";
+import { ModeledMethod } from "./modeled-method";
+import { Mode } from "./shared/mode";
+import { createDataExtensionYamls, loadDataExtensionYaml } from "./yaml";
+import { join } from "path";
+import { ExtensionPack } from "./shared/extension-pack";
+import {
+  Logger,
+  NotificationLogger,
+  showAndLogErrorMessage,
+} from "../common/logging";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { load as loadYaml } from "js-yaml";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { pathsEqual } from "../common/files";
+
+export async function saveModeledMethods(
+  extensionPack: ExtensionPack,
+  databaseName: string,
+  language: string,
+  externalApiUsages: ExternalApiUsage[],
+  modeledMethods: Record<string, ModeledMethod>,
+  mode: Mode,
+  logger: Logger,
+): Promise<void> {
+  const yamls = createDataExtensionYamls(
+    databaseName,
+    language,
+    externalApiUsages,
+    modeledMethods,
+    mode,
+  );
+
+  for (const [filename, yaml] of Object.entries(yamls)) {
+    await outputFile(join(extensionPack.path, filename), yaml);
+  }
+
+  void logger.log(`Saved data extension YAML`);
+}
+
+export async function loadModeledMethods(
+  extensionPack: ExtensionPack,
+  cliServer: CodeQLCliServer,
+  logger: NotificationLogger,
+): Promise<Record<string, ModeledMethod>> {
+  const modelFiles = await listModelFiles(extensionPack.path, cliServer);
+
+  const existingModeledMethods: Record<string, ModeledMethod> = {};
+
+  for (const modelFile of modelFiles) {
+    const yaml = await readFile(modelFile, "utf8");
+
+    const data = loadYaml(yaml, {
+      filename: modelFile,
+    });
+
+    const modeledMethods = loadDataExtensionYaml(data);
+    if (!modeledMethods) {
+      void showAndLogErrorMessage(
+        logger,
+        `Failed to parse data extension YAML ${modelFile}.`,
+      );
+      continue;
+    }
+
+    for (const [key, value] of Object.entries(modeledMethods)) {
+      existingModeledMethods[key] = value;
+    }
+  }
+
+  return existingModeledMethods;
+}
+
+export async function listModelFiles(
+  extensionPackPath: string,
+  cliServer: CodeQLCliServer,
+): Promise<Set<string>> {
+  const result = await cliServer.resolveExtensions(
+    extensionPackPath,
+    getOnDiskWorkspaceFolders(),
+  );
+
+  const modelFiles = new Set<string>();
+  for (const [path, extensions] of Object.entries(result.data)) {
+    if (pathsEqual(path, extensionPackPath)) {
+      for (const extension of extensions) {
+        modelFiles.add(extension.file);
+      }
+    }
+  }
+  return modelFiles;
+}
diff --git a/extensions/ql-vscode/src/data-extensions-editor/yaml.ts b/extensions/ql-vscode/src/data-extensions-editor/yaml.ts
@@ -9,6 +9,8 @@ import {
 
 import * as dataSchemaJson from "./data-schema.json";
 import { sanitizeExtensionPackName } from "./extension-pack-name";
+import { Mode } from "./shared/mode";
+import { assertNever } from "../common/helpers-pure";
 
 const ajv = new Ajv({ allErrors: true });
 const dataSchemaValidate = ajv.compile(dataSchemaJson);
@@ -66,6 +68,32 @@ export function createDataExtensionYaml(
 ${extensions.join("\n")}`;
 }
 
+export function createDataExtensionYamls(
+  databaseName: string,
+  language: string,
+  externalApiUsages: ExternalApiUsage[],
+  modeledMethods: Record<string, ModeledMethod>,
+  mode: Mode,
+) {
+  switch (mode) {
+    case Mode.Application:
+      return createDataExtensionYamlsForApplicationMode(
+        language,
+        externalApiUsages,
+        modeledMethods,
+      );
+    case Mode.Framework:
+      return createDataExtensionYamlsForFrameworkMode(
+        databaseName,
+        language,
+        externalApiUsages,
+        modeledMethods,
+      );
+    default:
+      assertNever(mode);
+  }
+}
+
 export function createDataExtensionYamlsForApplicationMode(
   language: string,
   externalApiUsages: ExternalApiUsage[],
diff --git a/extensions/ql-vscode/test/vscode-tests/cli-integration/data-extensions-editor/modeled-method-fs.test.ts b/extensions/ql-vscode/test/vscode-tests/cli-integration/data-extensions-editor/modeled-method-fs.test.ts
