Load existing data extension YAML in editor

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

@@ -15,17 +15,21 @@ import { extLogger, TeeLogger } from "../common";
 import { CoreCompletedQuery, QueryRunner } from "../queryRunner";
 import { qlpackOfDatabase } from "../contextual/queryResolver";
 import { file } from "tmp-promise";
-import { writeFile } from "fs-extra";
-import { dump } from "js-yaml";
+import { readFile, writeFile } from "fs-extra";
+import { dump as dumpYaml, load as loadYaml } from "js-yaml";
 import {
   getOnDiskWorkspaceFolders,
   showAndLogExceptionWithTelemetry,
+  showAndLogWarningMessage,
 } from "../helpers";
 import { DatabaseItem } from "../local-databases";
 import { CodeQLCliServer } from "../cli";
-import { assertNever, asError, getErrorMessage } from "../pure/helpers-pure";
+import { asError, assertNever, getErrorMessage } from "../pure/helpers-pure";
 import { decodeBqrsToExternalApiUsages } from "./bqrs";
 import { redactableError } from "../pure/errors";
+import { createDataExtensionYaml, loadDataExtensionYaml } from "./yaml";
+import { ExternalApiUsage } from "./external-api-usage";
+import { ModeledMethod } from "./modeled-method";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
@@ -70,8 +74,11 @@ export class DataExtensionsEditorView extends AbstractWebview<
         await this.onWebViewLoaded();
 
         break;
-      case "applyDataExtensionYaml":
-        await this.saveYaml(msg.yaml);
+      case "saveModeledMethods":
+        await this.saveModeledMethods(
+          msg.externalApiUsages,
+          msg.modeledMethods,
+        );
         await this.loadExternalApiUsages();
 
         break;
@@ -83,20 +90,57 @@ export class DataExtensionsEditorView extends AbstractWebview<
   protected async onWebViewLoaded() {
     super.onWebViewLoaded();
 
-    await this.loadExternalApiUsages();
+    await Promise.all([
+      this.loadExternalApiUsages(),
+      this.loadExistingModeledMethods(),
+    ]);
   }
 
-  protected async saveYaml(yaml: string): Promise<void> {
+  protected async saveModeledMethods(
+    externalApiUsages: ExternalApiUsage[],
+    modeledMethods: Record<string, ModeledMethod>,
+  ): Promise<void> {
     const modelFilename = this.calculateModelFilename();
     if (!modelFilename) {
       return;
     }
 
+    const yaml = createDataExtensionYaml(externalApiUsages, modeledMethods);
+
     await writeFile(modelFilename, yaml);
 
     void extLogger.log(`Saved data extension YAML to ${modelFilename}`);
   }
 
+  protected async loadExistingModeledMethods(): Promise<void> {
+    const modelFilename = this.calculateModelFilename();
+    if (!modelFilename) {
+      return;
+    }
+
+    try {
+      const yaml = await readFile(modelFilename, "utf8");
+
+      const data = loadYaml(yaml, {
+        filename: modelFilename,
+      });
+
+      const existingModeledMethods = loadDataExtensionYaml(data);
+
+      if (!existingModeledMethods) {
+        void showAndLogWarningMessage("Failed to parse data extension YAML.");
+        return;
+      }
+
+      await this.postMessage({
+        t: "setExistingModeledMethods",
+        existingModeledMethods,
+      });
+    } catch (e: unknown) {
+      void extLogger.log(`Unable to read data extension YAML: ${e}`);
+    }
+  }
+
   protected async loadExternalApiUsages(): Promise<void> {
     try {
       const queryResult = await this.runQuery();
@@ -165,7 +209,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
         },
       });
     }
-    await writeFile(suiteFile, dump(suiteYaml), "utf8");
+    await writeFile(suiteFile, dumpYaml(suiteYaml), "utf8");
 
     const queries = await this.cliServer.resolveQueriesInSuite(
       suiteFile,

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

@@ -9,8 +9,13 @@ type ExternalApiUsageByType = {
 type DataExtensionDefinition = {
   extensible: string;
   generateMethodDefinition: (method: ExternalApiUsageByType) => any[];
+  readModeledMethod: (row: any[]) => [string, ModeledMethod] | undefined;
 };
 
+function readRowToMethod(row: any[]): string {
+  return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
+}
+
 const definitions: Record<
   Exclude<ModeledMethodType, "none">,
   DataExtensionDefinition
@@ -32,6 +37,15 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      readRowToMethod(row),
+      {
+        type: "source",
+        input: "",
+        output: row[6],
+        kind: row[7],
+      },
+    ],
   },
   sink: {
     extensible: "sinkModel",
@@ -50,6 +64,15 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      readRowToMethod(row),
+      {
+        type: "sink",
+        input: row[6],
+        output: "",
+        kind: row[7],
+      },
+    ],
   },
   summary: {
     extensible: "summaryModel",
@@ -69,6 +92,15 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      readRowToMethod(row),
+      {
+        type: "summary",
+        input: row[6],
+        output: row[7],
+        kind: row[8],
+      },
+    ],
   },
   neutral: {
     extensible: "neutralModel",
@@ -82,6 +114,15 @@ const definitions: Record<
       method.externalApiUsage.methodParameters,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
+      {
+        type: "neutral",
+        input: "",
+        output: "",
+        kind: "",
+      },
+    ],
   },
 };
 
@@ -142,3 +183,55 @@ export function createDataExtensionYaml(
   return `extensions:
 ${extensions.join("\n")}`;
 }
+
+export function loadDataExtensionYaml(
+  data: any,
+): Record<string, ModeledMethod> | undefined {
+  if (typeof data !== "object") {
+    return undefined;
+  }
+
+  const extensions = data.extensions;
+  if (!Array.isArray(extensions)) {
+    return undefined;
+  }
+
+  const modeledMethods: Record<string, ModeledMethod> = {};
+
+  for (const extension of extensions) {
+    const addsTo = extension.addsTo;
+    if (typeof addsTo !== "object") {
+      continue;
+    }
+
+    const extensible = addsTo.extensible;
+    if (typeof extensible !== "string") {
+      continue;
+    }
+
+    const data = extension.data;
+    if (!Array.isArray(data)) {
+      continue;
+    }
+
+    const definition = Object.values(definitions).find(
+      (definition) => definition.extensible === extensible,
+    );
+    if (!definition) {
+      continue;
+    }
+
+    for (const row of data) {
+      const result = definition.readModeledMethod(row);
+      if (!result) {
+        continue;
+      }
+
+      const [apiInfo, modeledMethod] = result;
+
+      modeledMethods[apiInfo] = modeledMethod;
+    }
+  }
+
+  return modeledMethods;
+}

extensions/ql-vscode/src/pure/interface-types.ts

@@ -15,6 +15,7 @@ import { RepositoriesFilterSortStateWithIds } from "./variant-analysis-filter-so
 import { ErrorLike } from "./errors";
 import { DataFlowPaths } from "../variant-analysis/shared/data-flow-paths";
 import { ExternalApiUsage } from "../data-extensions-editor/external-api-usage";
+import { ModeledMethod } from "../data-extensions-editor/modeled-method";
 
 /**
  * This module contains types and code that are shared between
@@ -492,15 +493,22 @@ export interface ShowProgressMessage {
   message: string;
 }
 
-export interface ApplyDataExtensionYamlMessage {
-  t: "applyDataExtensionYaml";
-  yaml: string;
+export interface SetExistingModeledMethods {
+  t: "setExistingModeledMethods";
+  existingModeledMethods: Record<string, ModeledMethod>;
+}
+
+export interface SaveModeledMethods {
+  t: "saveModeledMethods";
+  externalApiUsages: ExternalApiUsage[];
+  modeledMethods: Record<string, ModeledMethod>;
 }
 
 export type ToDataExtensionsEditorMessage =
   | SetExternalApiUsagesMessage
-  | ShowProgressMessage;
+  | ShowProgressMessage
+  | SetExistingModeledMethods;
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg
-  | ApplyDataExtensionYamlMessage;
+  | SaveModeledMethods;

extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx

@@ -16,7 +16,6 @@ import { ModeledMethod } from "../../data-extensions-editor/modeled-method";
 import { MethodRow } from "./MethodRow";
 import { assertNever } from "../../pure/helpers-pure";
 import { vscode } from "../vscode-api";
-import { createDataExtensionYaml } from "../../data-extensions-editor/yaml";
 import { calculateSupportedPercentage } from "./supported";
 
 export const DataExtensionsEditorContainer = styled.div`
@@ -57,6 +56,15 @@ export function DataExtensionsEditor(): JSX.Element {
             break;
           case "showProgress":
             setProgress(msg);
+            break;
+          case "setExistingModeledMethods":
+            setModeledMethods((oldModeledMethods) => {
+              return {
+                ...msg.existingModeledMethods,
+                ...oldModeledMethods,
+              };
+            });
+
             break;
           default:
             assertNever(msg);
@@ -92,14 +100,10 @@ export function DataExtensionsEditor(): JSX.Element {
   );
 
   const onApplyClick = useCallback(() => {
-    const yamlString = createDataExtensionYaml(
+    vscode.postMessage({
+      t: "saveModeledMethods",
       externalApiUsages,
       modeledMethods,
-    );
-
-    vscode.postMessage({
-      t: "applyDataExtensionYaml",
-      yaml: yamlString,
     });
   }, [externalApiUsages, modeledMethods]);
 

extensions/ql-vscode/test/unit-tests/data-extensions-editor/yaml.test.ts

@@ -1,4 +1,7 @@
-import { createDataExtensionYaml } from "../../../src/data-extensions-editor/yaml";
+import {
+  createDataExtensionYaml,
+  loadDataExtensionYaml,
+} from "../../../src/data-extensions-editor/yaml";
 
 describe("createDataExtensionYaml", () => {
   it("creates the correct YAML file", () => {
@@ -99,3 +102,61 @@ describe("createDataExtensionYaml", () => {
 `);
   });
 });
+
+describe("loadDataExtensionYaml", () => {
+  it("loads the YAML file", () => {
+    const data = loadDataExtensionYaml({
+      extensions: [
+        {
+          addsTo: { pack: "codeql/java-all", extensible: "sourceModel" },
+          data: [],
+        },
+        {
+          addsTo: { pack: "codeql/java-all", extensible: "sinkModel" },
+          data: [
+            [
+              "org.sql2o",
+              "Connection",
+              true,
+              "createQuery",
+              "(String)",
+              "",
+              "Argument[0]",
+              "sql",
+              "manual",
+            ],
+          ],
+        },
+        {
+          addsTo: { pack: "codeql/java-all", extensible: "summaryModel" },
+          data: [],
+        },
+        {
+          addsTo: { pack: "codeql/java-all", extensible: "neutralModel" },
+          data: [],
+        },
+      ],
+    });
+
+    expect(data).toEqual({
+      "org.sql2o.Connection#createQuery(String)": {
+        input: "Argument[0]",
+        kind: "sql",
+        output: "",
+        type: "sink",
+      },
+    });
+  });
+
+  it("returns undefined if given a string", () => {
+    const data = loadDataExtensionYaml(`extensions:
+  - addsTo:
+      pack: codeql/java-all
+      extensible: sinkModel
+    data:
+      - ["org.sql2o","Connection",true,"createQuery","(String)","","Argument[0]","sql","manual"]
+`);
+
+    expect(data).toBeUndefined();
+  });
+});