Generate separate file for generated type models in Ruby

Author: koesie10

extensions/ql-vscode/src/model-editor/generate.ts

@@ -5,19 +5,15 @@ import type { QueryRunner } from "../query-server";
 import type { CodeQLCliServer } from "../codeql-cli/cli";
 import type { ProgressCallback } from "../common/vscode/progress";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
-import type { ModeledMethod } from "./modeled-method";
 import { runQuery } from "../local-queries/run-query";
 import type { QueryConstraints } from "../local-queries";
 import { resolveQueries } from "../local-queries";
 import type { DecodedBqrs } from "../common/bqrs-cli-types";
+
 type GenerateQueriesOptions = {
   queryConstraints: QueryConstraints;
   filterQueries?: (queryPath: string) => boolean;
-  parseResults: (
-    queryPath: string,
-    results: DecodedBqrs,
-  ) => ModeledMethod[] | Promise<ModeledMethod[]>;
-  onResults: (results: ModeledMethod[]) => void | Promise<void>;
+  onResults: (queryPath: string, results: DecodedBqrs) => void | Promise<void>;
 
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
@@ -28,7 +24,7 @@ type GenerateQueriesOptions = {
 };
 
 export async function runGenerateQueries(options: GenerateQueriesOptions) {
-  const { queryConstraints, filterQueries, parseResults, onResults } = options;
+  const { queryConstraints, filterQueries, onResults } = options;
 
   options.progress({
     message: "Resolving queries",
@@ -55,7 +51,7 @@ export async function runGenerateQueries(options: GenerateQueriesOptions) {
 
     const bqrs = await runSingleGenerateQuery(queryPath, i, maxStep, options);
     if (bqrs) {
-      await onResults(await parseResults(queryPath, bqrs));
+      await onResults(queryPath, bqrs);
     }
   }
 }

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -7,7 +7,7 @@ import type {
   SummaryModeledMethod,
   TypeModeledMethod,
 } from "../modeled-method";
-import type { DataTuple } from "../model-extension-file";
+import type { DataTuple, ModelExtension } from "../model-extension-file";
 import type { Mode } from "../shared/mode";
 import type { QueryConstraints } from "../../local-queries/query-constraints";
 import type {
@@ -32,6 +32,11 @@ export type ModelsAsDataLanguagePredicate<T> = {
   readModeledMethod: ReadModeledMethod;
 };
 
+export type GenerationContext = {
+  mode: Mode;
+  isCanary: boolean;
+};
+
 type ParseGenerationResults = (
   // The path to the query that generated the results.
   queryPath: string,
@@ -42,24 +47,37 @@ type ParseGenerationResults = (
   modelsAsDataLanguage: ModelsAsDataLanguage,
   // The logger to use for logging.
   logger: BaseLogger,
+  // Context about this invocation of the generation.
+  context: GenerationContext,
 ) => ModeledMethod[];
 
 type ModelsAsDataLanguageModelGeneration = {
   queryConstraints: (mode: Mode) => QueryConstraints;
   filterQueries?: (queryPath: string) => boolean;
   parseResults: ParseGenerationResults;
+};
+
+type ParseResultsToYaml = (
+  // The path to the query that generated the results.
+  queryPath: string,
+  // The results of the query.
+  bqrs: DecodedBqrs,
+  // The language-specific predicate that was used to generate the results. This is passed to allow
+  // sharing of code between different languages.
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  // The logger to use for logging.
+  logger: BaseLogger,
+) => ModelExtension[];
+
+type ModelsAsDataLanguageAutoModelGeneration = {
+  queryConstraints: (mode: Mode) => QueryConstraints;
+  filterQueries?: (queryPath: string) => boolean;
+  parseResultsToYaml: ParseResultsToYaml;
   /**
-   * If autoRun is not undefined, the query will be run automatically when the user starts the
-   * model editor.
-   *
-   * This only applies to framework mode. Application mode will never run the query automatically.
+   * By default, auto model generation is enabled for all modes. This function can be used to
+   * override that behavior.
    */
-  autoRun?: {
-    /**
-     * If defined, will use a custom parsing function when the query is run automatically.
-     */
-    parseResults?: ParseGenerationResults;
-  };
+  enabled?: (context: GenerationContext) => boolean;
 };
 
 type ModelsAsDataLanguageAccessPathSuggestions = {
@@ -109,6 +127,7 @@ export type ModelsAsDataLanguage = {
   ) => EndpointType | undefined;
   predicates: ModelsAsDataLanguagePredicates;
   modelGeneration?: ModelsAsDataLanguageModelGeneration;
+  autoModelGeneration?: ModelsAsDataLanguageAutoModelGeneration;
   accessPathSuggestions?: ModelsAsDataLanguageAccessPathSuggestions;
   /**
    * Returns the list of valid arguments that can be selected for the given method.

extensions/ql-vscode/src/model-editor/languages/ruby/generate.ts

@@ -1,6 +1,9 @@
 import type { BaseLogger } from "../../../common/logging";
 import type { DecodedBqrs } from "../../../common/bqrs-cli-types";
-import type { ModelsAsDataLanguage } from "../models-as-data";
+import type {
+  GenerationContext,
+  ModelsAsDataLanguage,
+} from "../models-as-data";
 import type { ModeledMethod } from "../../modeled-method";
 import type { DataTuple } from "../../model-extension-file";
 
@@ -9,10 +12,21 @@ export function parseGenerateModelResults(
   bqrs: DecodedBqrs,
   modelsAsDataLanguage: ModelsAsDataLanguage,
   logger: BaseLogger,
+  { isCanary }: GenerationContext,
 ): ModeledMethod[] {
   const modeledMethods: ModeledMethod[] = [];
 
   for (const resultSetName in bqrs) {
+    if (
+      resultSetName ===
+        modelsAsDataLanguage.predicates.type?.extensiblePredicate &&
+      !isCanary
+    ) {
+      // Don't load generated type results in non-canary mode. These are already automatically
+      // generated on start-up.
+      continue;
+    }
+
     const definition = Object.values(modelsAsDataLanguage.predicates).find(
       (definition) => definition.extensiblePredicate === resultSetName,
     );

extensions/ql-vscode/src/model-editor/languages/ruby/index.ts

@@ -177,28 +177,39 @@ export const ruby: ModelsAsDataLanguage = {
       "tags contain all": ["modeleditor", "generate-model", modeTag(mode)],
     }),
     parseResults: parseGenerateModelResults,
-    autoRun: {
-      parseResults: (queryPath, bqrs, modelsAsDataLanguage, logger) => {
-        // Only type models are generated automatically
-        const typePredicate = modelsAsDataLanguage.predicates.type;
-        if (!typePredicate) {
-          throw new Error("Type predicate not found");
-        }
+  },
+  autoModelGeneration: {
+    queryConstraints: (mode) => ({
+      kind: "table",
+      "tags contain all": ["modeleditor", "generate-model", modeTag(mode)],
+    }),
+    parseResultsToYaml: (_queryPath, bqrs, modelsAsDataLanguage) => {
+      const typePredicate = modelsAsDataLanguage.predicates.type;
+      if (!typePredicate) {
+        throw new Error("Type predicate not found");
+      }
 
-        const filteredBqrs = Object.fromEntries(
-          Object.entries(bqrs).filter(
-            ([key]) => key === typePredicate.extensiblePredicate,
-          ),
-        );
+      const typeTuples = bqrs[typePredicate.extensiblePredicate];
+      if (!typeTuples) {
+        return [];
+      }
 
-        return parseGenerateModelResults(
-          queryPath,
-          filteredBqrs,
-          modelsAsDataLanguage,
-          logger,
-        );
-      },
+      return [
+        {
+          addsTo: {
+            pack: "codeql/ruby-all",
+            extensible: typePredicate.extensiblePredicate,
+          },
+          data: typeTuples.tuples.filter((tuple): tuple is string[] => {
+            return (
+              tuple.filter((x) => typeof x === "string").length === tuple.length
+            );
+          }),
+        },
+      ];
     },
+    // Only enabled for framework mode in non-canary
+    enabled: ({ mode, isCanary }) => mode === Mode.Framework && !isCanary,
   },
   accessPathSuggestions: {
     queryConstraints: (mode) => ({

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

@@ -40,8 +40,13 @@ import type { Method } from "./method";
 import type { ModeledMethod } from "./modeled-method";
 import type { ExtensionPack } from "./shared/extension-pack";
 import type { ModelConfigListener } from "../config";
+import { isCanary } from "../config";
 import { Mode } from "./shared/mode";
-import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
+import {
+  GENERATED_MODELS_SUFFIX,
+  loadModeledMethods,
+  saveModeledMethods,
+} from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
 import type { QueryLanguage } from "../common/query-language";
 import { getLanguageDisplayName } from "../common/query-language";
@@ -60,6 +65,10 @@ import { parseAccessPathSuggestionRowsToOptions } from "./suggestions-bqrs";
 import { ModelEvaluator } from "./model-evaluator";
 import type { ModelEvaluationRunState } from "./shared/model-evaluation-run-state";
 import type { VariantAnalysisManager } from "../variant-analysis/variant-analysis-manager";
+import type { ModelExtensionFile } from "./model-extension-file";
+import { modelExtensionFileToYaml } from "./yaml";
+import { outputFile } from "fs-extra";
+import { join } from "path";
 
 export class ModelEditorView extends AbstractWebview<
   ToModelEditorMessage,
@@ -645,14 +654,18 @@ export class ModelEditorView extends AbstractWebview<
           await runGenerateQueries({
             queryConstraints: modelGeneration.queryConstraints(mode),
             filterQueries: modelGeneration.filterQueries,
-            parseResults: (queryPath, results) =>
-              modelGeneration.parseResults(
+            onResults: async (queryPath, results) => {
+              const modeledMethods = modelGeneration.parseResults(
                 queryPath,
                 results,
                 modelsAsDataLanguage,
                 this.app.logger,
-              ),
-            onResults: async (modeledMethods) => {
+                {
+                  mode,
+                  isCanary: isCanary(),
+                },
+              );
+
               this.addModeledMethodsFromArray(modeledMethods);
             },
             cliServer: this.cliServer,
@@ -678,15 +691,17 @@ export class ModelEditorView extends AbstractWebview<
 
   protected async generateModeledMethodsOnStartup(): Promise<void> {
     const mode = this.modelingStore.getMode(this.databaseItem);
-    if (mode !== Mode.Framework) {
+    const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
+    const autoModelGeneration = modelsAsDataLanguage.autoModelGeneration;
+
+    if (autoModelGeneration === undefined) {
       return;
     }
 
-    const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
-    const modelGeneration = modelsAsDataLanguage.modelGeneration;
-    const autoRun = modelGeneration?.autoRun;
-
-    if (modelGeneration === undefined || autoRun === undefined) {
+    if (
+      autoModelGeneration.enabled &&
+      !autoModelGeneration.enabled({ mode, isCanary: isCanary() })
+    ) {
       return;
     }
 
@@ -698,22 +713,23 @@ export class ModelEditorView extends AbstractWebview<
           message: "Generating models",
         });
 
-        const parseResults =
-          autoRun.parseResults ?? modelGeneration.parseResults;
+        const extensionFile: ModelExtensionFile = {
+          extensions: [],
+        };
 
         try {
           await runGenerateQueries({
-            queryConstraints: modelGeneration.queryConstraints(mode),
-            filterQueries: modelGeneration.filterQueries,
-            parseResults: (queryPath, results) =>
-              parseResults(
+            queryConstraints: autoModelGeneration.queryConstraints(mode),
+            filterQueries: autoModelGeneration.filterQueries,
+            onResults: (queryPath, results) => {
+              const extensions = autoModelGeneration.parseResultsToYaml(
                 queryPath,
                 results,
                 modelsAsDataLanguage,
                 this.app.logger,
-              ),
-            onResults: async (modeledMethods) => {
-              this.addModeledMethodsFromArray(modeledMethods);
+              );
+
+              extensionFile.extensions.push(...extensions);
             },
             cliServer: this.cliServer,
             queryRunner: this.queryRunner,
@@ -730,7 +746,25 @@ export class ModelEditorView extends AbstractWebview<
               asError(e),
             )`Failed to auto-run generating models: ${getErrorMessage(e)}`,
           );
+          return;
         }
+
+        progress({
+          step: 4000,
+          maxStep: 4000,
+          message: "Saving generated models",
+        });
+
+        const fileContents = `# This file was automatically generated based from ${this.databaseItem.name}. Manual changes will not persist.\n\n${modelExtensionFileToYaml(extensionFile)}`;
+        const filePath = join(
+          this.extensionPack.path,
+          "models",
+          `${this.language}${GENERATED_MODELS_SUFFIX}`,
+        );
+
+        await outputFile(filePath, fileContents);
+
+        void this.app.logger.log(`Saved generated model file to ${filePath}`);
       },
       {
         cancellable: false,

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

@@ -12,6 +12,9 @@ import { load as loadYaml } from "js-yaml";
 import type { CodeQLCliServer } from "../codeql-cli/cli";
 import { pathsEqual } from "../common/files";
 import type { QueryLanguage } from "../common/query-language";
+import { isCanary } from "../config";
+
+export const GENERATED_MODELS_SUFFIX = ".model.generated.yml";
 
 export async function saveModeledMethods(
   extensionPack: ExtensionPack,
@@ -118,6 +121,11 @@ export async function listModelFiles(
   for (const [path, extensions] of Object.entries(result.data)) {
     if (pathsEqual(path, extensionPackPath)) {
       for (const extension of extensions) {
+        // We only load generated models in canary mode
+        if (!isCanary() && extension.file.endsWith(GENERATED_MODELS_SUFFIX)) {
+          continue;
+        }
+
         modelFiles.add(relative(extensionPackPath, extension.file));
       }
     }

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

@@ -337,7 +337,7 @@ function validateModelExtensionFile(data: unknown): data is ModelExtensionFile {
  *
  * @param data The data extension file
  */
-function modelExtensionFileToYaml(data: ModelExtensionFile) {
+export function modelExtensionFileToYaml(data: ModelExtensionFile) {
   const extensions = data.extensions
     .map((extension) => {
       const data =

extensions/ql-vscode/test/unit-tests/model-editor/languages/ruby/generate.test.ts

@@ -4,6 +4,7 @@ import { ruby } from "../../../../../src/model-editor/languages/ruby";
 import { createMockLogger } from "../../../../__mocks__/loggerMock";
 import type { ModeledMethod } from "../../../../../src/model-editor/modeled-method";
 import { EndpointType } from "../../../../../src/model-editor/method";
+import { Mode } from "../../../../../src/model-editor/shared/mode";
 
 describe("parseGenerateModelResults", () => {
   it("should return the results", async () => {
@@ -76,6 +77,10 @@ describe("parseGenerateModelResults", () => {
       bqrs,
       ruby,
       createMockLogger(),
+      {
+        isCanary: true,
+        mode: Mode.Framework,
+      },
     );
     expect(result).toEqual([
       {