Load existing data extension YAML in editor

This loads in the existing data extension YAML file for the selected
database. It only supports the filename we save it to, and will not load
it from any other data extension YAML files.

Author: koesie10

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

@@ -15,8 +15,8 @@ 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, load } from "js-yaml";
 import { getOnDiskWorkspaceFolders } from "../helpers";
 import { DatabaseItem } from "../local-databases";
 import { CodeQLCliServer } from "../cli";
@@ -78,7 +78,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
   protected async onWebViewLoaded() {
     super.onWebViewLoaded();
 
-    await this.loadExternalApiUsages();
+    await Promise.all([this.loadExternalApiUsages(), this.readExistingYaml()]);
   }
 
   protected async saveYaml(yaml: string): Promise<void> {
@@ -92,6 +92,28 @@ export class DataExtensionsEditorView extends AbstractWebview<
     void extLogger.log(`Saved data extension YAML to ${modelFilename}`);
   }
 
+  protected async readExistingYaml(): Promise<void> {
+    const modelFilename = this.modelFileName;
+    if (!modelFilename) {
+      return;
+    }
+
+    try {
+      const yaml = await readFile(modelFilename, "utf8");
+
+      const data = load(yaml, {
+        filename: modelFilename,
+      });
+
+      await this.postMessage({
+        t: "setExistingYamlData",
+        data,
+      });
+    } catch (e: unknown) {
+      void extLogger.log(`Unable to read data extension YAML: ${e}`);
+    }
+  }
+
   protected async loadExternalApiUsages(): Promise<void> {
     const queryResult = await this.runQuery();
     if (!queryResult) {

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

@@ -12,8 +12,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
@@ -35,6 +40,15 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      readRowToMethod(row),
+      {
+        type: "source",
+        input: "",
+        output: row[6],
+        kind: row[7],
+      },
+    ],
   },
   sink: {
     extensible: "sinkModel",
@@ -53,6 +67,15 @@ const definitions: Record<
       method.modeledMethod.kind,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      readRowToMethod(row),
+      {
+        type: "sink",
+        input: row[6],
+        output: "",
+        kind: row[7],
+      },
+    ],
   },
   summary: {
     extensible: "summaryModel",
@@ -72,6 +95,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",
@@ -85,6 +117,15 @@ const definitions: Record<
       method.method.methodParameters,
       "manual",
     ],
+    readModeledMethod: (row) => [
+      `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
+      {
+        type: "neutral",
+        input: "",
+        output: "",
+        kind: "",
+      },
+    ],
   },
 };
 
@@ -145,3 +186,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

@@ -492,14 +492,20 @@ export interface ShowProgressMessage {
   message: string;
 }
 
+export interface SetExistingYamlDataMessage {
+  t: "setExistingYamlData";
+  data: any;
+}
+
 export interface ApplyDataExtensionYamlMessage {
   t: "applyDataExtensionYaml";
   yaml: string;
 }
 
 export type ToDataExtensionsEditorMessage =
   | SetExternalApiResultsMessage
-  | ShowProgressMessage;
+  | ShowProgressMessage
+  | SetExistingYamlDataMessage;
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg

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

@@ -20,7 +20,10 @@ import {
 import { MethodRow } from "./MethodRow";
 import { assertNever } from "../../pure/helpers-pure";
 import { vscode } from "../vscode-api";
-import { createDataExtensionYaml } from "../../data-extensions-editor/yaml";
+import {
+  createDataExtensionYaml,
+  loadDataExtensionYaml,
+} from "../../data-extensions-editor/yaml";
 
 export const DataExtensionsEditorContainer = styled.div`
   margin-top: 1rem;
@@ -60,6 +63,17 @@ export function DataExtensionsEditor(): JSX.Element {
             break;
           case "showProgress":
             setProgress(msg);
+            break;
+          case "setExistingYamlData":
+            setModeledMethods((oldModeledMethods) => {
+              const existingModeledMethods = loadDataExtensionYaml(msg.data);
+
+              return {
+                ...existingModeledMethods,
+                ...oldModeledMethods,
+              };
+            });
+
             break;
           default:
             assertNever(msg);

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();
+  });
+});