Load existing modeled methods before saving, to avoid overwriting data

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

@@ -123,6 +123,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
           msg.externalApiUsages,
           msg.modeledMethods,
           this.mode,
+          this.cliServer,
           this.app.logger,
         );
         await Promise.all([this.setViewState(), this.loadExternalApiUsages()]);

extensions/ql-vscode/src/data-extensions-editor/modeled-method-fs.ts

@@ -3,13 +3,9 @@ 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 { join, relative } from "path";
 import { ExtensionPack } from "./shared/extension-pack";
-import {
-  Logger,
-  NotificationLogger,
-  showAndLogErrorMessage,
-} from "../common/logging";
+import { 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";
@@ -22,13 +18,21 @@ export async function saveModeledMethods(
   externalApiUsages: ExternalApiUsage[],
   modeledMethods: Record<string, ModeledMethod>,
   mode: Mode,
-  logger: Logger,
+  cliServer: CodeQLCliServer,
+  logger: NotificationLogger,
 ): Promise<void> {
+  const existingModeledMethods = await loadModeledMethodFiles(
+    extensionPack,
+    cliServer,
+    logger,
+  );
+
   const yamls = createDataExtensionYamls(
     databaseName,
     language,
     externalApiUsages,
     modeledMethods,
+    existingModeledMethods,
     mode,
   );
 
@@ -39,17 +43,20 @@ export async function saveModeledMethods(
   void logger.log(`Saved data extension YAML`);
 }
 
-export async function loadModeledMethods(
+async function loadModeledMethodFiles(
   extensionPack: ExtensionPack,
   cliServer: CodeQLCliServer,
   logger: NotificationLogger,
-): Promise<Record<string, ModeledMethod>> {
+): Promise<Record<string, Record<string, ModeledMethod>>> {
   const modelFiles = await listModelFiles(extensionPack.path, cliServer);
 
-  const existingModeledMethods: Record<string, ModeledMethod> = {};
+  const modeledMethodsByFile: Record<
+    string,
+    Record<string, ModeledMethod>
+  > = {};
 
   for (const modelFile of modelFiles) {
-    const yaml = await readFile(modelFile, "utf8");
+    const yaml = await readFile(join(extensionPack.path, modelFile), "utf8");
 
     const data = loadYaml(yaml, {
       filename: modelFile,
@@ -63,7 +70,25 @@ export async function loadModeledMethods(
       );
       continue;
     }
+    modeledMethodsByFile[modelFile] = modeledMethods;
+  }
+
+  return modeledMethodsByFile;
+}
 
+export async function loadModeledMethods(
+  extensionPack: ExtensionPack,
+  cliServer: CodeQLCliServer,
+  logger: NotificationLogger,
+): Promise<Record<string, ModeledMethod>> {
+  const existingModeledMethods: Record<string, ModeledMethod> = {};
+
+  const modeledMethodsByFile = await loadModeledMethodFiles(
+    extensionPack,
+    cliServer,
+    logger,
+  );
+  for (const modeledMethods of Object.values(modeledMethodsByFile)) {
     for (const [key, value] of Object.entries(modeledMethods)) {
       existingModeledMethods[key] = value;
     }
@@ -85,7 +110,7 @@ export async function listModelFiles(
   for (const [path, extensions] of Object.entries(result.data)) {
     if (pathsEqual(path, extensionPackPath)) {
       for (const extension of extensions) {
-        modelFiles.add(extension.file);
+        modelFiles.add(relative(extensionPackPath, extension.file));
       }
     }
   }

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

@@ -72,22 +72,25 @@ export function createDataExtensionYamls(
   databaseName: string,
   language: string,
   externalApiUsages: ExternalApiUsage[],
-  modeledMethods: Record<string, ModeledMethod>,
+  newModeledMethods: Record<string, ModeledMethod>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
   mode: Mode,
 ) {
   switch (mode) {
     case Mode.Application:
       return createDataExtensionYamlsForApplicationMode(
         language,
         externalApiUsages,
-        modeledMethods,
+        newModeledMethods,
+        existingModeledMethods,
       );
     case Mode.Framework:
       return createDataExtensionYamlsForFrameworkMode(
         databaseName,
         language,
         externalApiUsages,
-        modeledMethods,
+        newModeledMethods,
+        existingModeledMethods,
       );
     default:
       assertNever(mode);
@@ -97,27 +100,51 @@ export function createDataExtensionYamls(
 export function createDataExtensionYamlsForApplicationMode(
   language: string,
   externalApiUsages: ExternalApiUsage[],
-  modeledMethods: Record<string, ModeledMethod>,
+  newModeledMethods: Record<string, ModeledMethod>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
 ): Record<string, string> {
-  const methodsByLibraryFilename: Record<string, ModeledMethod[]> = {};
-
+  const methodsByLibraryFilename: Record<
+    string,
+    Record<string, ModeledMethod>
+  > = {};
+
+  // We only want to generate a yaml file when it's a known external API usage
+  // and there are new modeled methods for it. This avoids us overwriting other
+  // files that may contain data we don't know about.
   for (const externalApiUsage of externalApiUsages) {
-    const modeledMethod = modeledMethods[externalApiUsage.signature];
-    if (!modeledMethod) {
-      continue;
+    if (externalApiUsage.signature in newModeledMethods) {
+      methodsByLibraryFilename[
+        createFilenameForLibrary(externalApiUsage.library)
+      ] = {};
     }
+  }
 
-    const filename = createFilenameForLibrary(externalApiUsage.library);
+  // First populate methodsByLibraryFilename with any existing modeled methods.
+  for (const [filename, methods] of Object.entries(existingModeledMethods)) {
+    if (filename in methodsByLibraryFilename) {
+      for (const [signature, method] of Object.entries(methods)) {
+        methodsByLibraryFilename[filename][signature] = method;
+      }
+    }
+  }
 
-    methodsByLibraryFilename[filename] =
-      methodsByLibraryFilename[filename] || [];
-    methodsByLibraryFilename[filename].push(modeledMethod);
+  // Add the new modeled methods, potentially overwriting existing modeled methods
+  // but not removing existing modeled methods that are not in the new set.
+  for (const externalApiUsage of externalApiUsages) {
+    const method = newModeledMethods[externalApiUsage.signature];
+    if (method) {
+      const filename = createFilenameForLibrary(externalApiUsage.library);
+      methodsByLibraryFilename[filename][method.signature] = method;
+    }
   }
 
   const result: Record<string, string> = {};
 
   for (const [filename, methods] of Object.entries(methodsByLibraryFilename)) {
-    result[filename] = createDataExtensionYaml(language, methods);
+    result[filename] = createDataExtensionYaml(
+      language,
+      Object.values(methods),
+    );
   }
 
   return result;
@@ -127,7 +154,8 @@ export function createDataExtensionYamlsForFrameworkMode(
   databaseName: string,
   language: string,
   externalApiUsages: ExternalApiUsage[],
-  modeledMethods: Record<string, ModeledMethod>,
+  newModeledMethods: Record<string, ModeledMethod>,
+  existingModeledMethods: Record<string, Record<string, ModeledMethod>>,
   prefix = "models/",
   suffix = ".model",
 ): Record<string, string> {
@@ -136,16 +164,28 @@ export function createDataExtensionYamlsForFrameworkMode(
     .slice(1)
     .map((part) => sanitizeExtensionPackName(part))
     .join("-");
+  const filename = `${prefix}${libraryName}${suffix}.yml`;
+
+  const methods: Record<string, ModeledMethod> = {};
+
+  // First populate methodsByLibraryFilename with any existing modeled methods.
+  for (const [signature, method] of Object.entries(
+    existingModeledMethods[filename] || {},
+  )) {
+    methods[signature] = method;
+  }
 
-  const methods = externalApiUsages
-    .map((externalApiUsage) => modeledMethods[externalApiUsage.signature])
-    .filter((modeledMethod) => modeledMethod !== undefined);
+  // Add the new modeled methods, potentially overwriting existing modeled methods
+  // but not removing existing modeled methods that are not in the new set.
+  for (const externalApiUsage of externalApiUsages) {
+    const modeledMethod = newModeledMethods[externalApiUsage.signature];
+    if (modeledMethod) {
+      methods[modeledMethod.signature] = modeledMethod;
+    }
+  }
 
   return {
-    [`${prefix}${libraryName}${suffix}.yml`]: createDataExtensionYaml(
-      language,
-      methods,
-    ),
+    [filename]: createDataExtensionYaml(language, Object.values(methods)),
   };
 }