Auto-generate type models for Ruby

Author: koesie10

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

@@ -32,20 +32,34 @@ export type ModelsAsDataLanguagePredicate<T> = {
   readModeledMethod: ReadModeledMethod;
 };
 
+type ParseGenerationResults = (
+  // 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,
+) => ModeledMethod[];
+
 type ModelsAsDataLanguageModelGeneration = {
   queryConstraints: (mode: Mode) => QueryConstraints;
   filterQueries?: (queryPath: string) => boolean;
-  parseResults: (
-    // 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,
-  ) => ModeledMethod[];
+  parseResults: ParseGenerationResults;
+  /**
+   * 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.
+   */
+  autoRun?: {
+    /**
+     * If defined, will use a custom parsing function when the query is run automatically.
+     */
+    parseResults?: ParseGenerationResults;
+  };
 };
 
 type ModelsAsDataLanguageAccessPathSuggestions = {

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

@@ -177,6 +177,28 @@ 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");
+        }
+
+        const filteredBqrs = Object.fromEntries(
+          Object.entries(bqrs).filter(
+            ([key]) => key === typePredicate.extensiblePredicate,
+          ),
+        );
+
+        return parseGenerateModelResults(
+          queryPath,
+          filteredBqrs,
+          modelsAsDataLanguage,
+          logger,
+        );
+      },
+    },
   },
   accessPathSuggestions: {
     queryConstraints: (mode) => ({

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

@@ -64,6 +64,9 @@ export class ModelEditorView extends AbstractWebview<
 > {
   private readonly autoModeler: AutoModeler;
   private readonly languageDefinition: ModelsAsDataLanguage;
+  // Cancellation token source that can be used for passing into long-running operations. Should only
+  // be cancelled when the view is closed
+  private readonly cancellationTokenSource = new CancellationTokenSource();
 
   public constructor(
     protected readonly app: App,
@@ -83,6 +86,12 @@ export class ModelEditorView extends AbstractWebview<
   ) {
     super(app);
 
+    this.push({
+      dispose: () => {
+        this.cancellationTokenSource.cancel();
+      },
+    });
+
     this.modelingStore.initializeStateForDb(databaseItem, initialMode);
     this.registerToModelingEvents();
     this.registerToModelConfigEvents();
@@ -367,6 +376,8 @@ export class ModelEditorView extends AbstractWebview<
       this.setViewState(),
       withProgress((progress, token) => this.loadMethods(progress, token), {
         cancellable: true,
+      }).then(async () => {
+        await this.generateModeledMethodsOnStartup();
       }),
       this.loadExistingModeledMethods(),
       // Only load access path suggestions if the feature is enabled
@@ -471,7 +482,7 @@ export class ModelEditorView extends AbstractWebview<
 
     try {
       if (!token) {
-        token = new CancellationTokenSource().token;
+        token = this.cancellationTokenSource.token;
       }
       const queryResult = await runModelEditorQueries(mode, {
         cliServer: this.cliServer,
@@ -522,8 +533,6 @@ export class ModelEditorView extends AbstractWebview<
   protected async loadAccessPathSuggestions(
     progress: ProgressCallback,
   ): Promise<void> {
-    const tokenSource = new CancellationTokenSource();
-
     const mode = this.modelingStore.getMode(this.databaseItem);
 
     const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
@@ -546,7 +555,7 @@ export class ModelEditorView extends AbstractWebview<
         queryStorageDir: this.queryStorageDir,
         databaseItem: this.databaseItem,
         progress,
-        token: tokenSource.token,
+        token: this.cancellationTokenSource.token,
         logger: this.app.logger,
       });
 
@@ -577,8 +586,6 @@ export class ModelEditorView extends AbstractWebview<
   protected async generateModeledMethods(): Promise<void> {
     await withProgress(
       async (progress) => {
-        const tokenSource = new CancellationTokenSource();
-
         const mode = this.modelingStore.getMode(this.databaseItem);
 
         const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
@@ -636,7 +643,7 @@ export class ModelEditorView extends AbstractWebview<
             queryStorageDir: this.queryStorageDir,
             databaseItem: addedDatabase ?? this.databaseItem,
             progress,
-            token: tokenSource.token,
+            token: this.cancellationTokenSource.token,
           });
         } catch (e: unknown) {
           void showAndLogExceptionWithTelemetry(
@@ -652,6 +659,70 @@ export class ModelEditorView extends AbstractWebview<
     );
   }
 
+  protected async generateModeledMethodsOnStartup(): Promise<void> {
+    const mode = this.modelingStore.getMode(this.databaseItem);
+    if (mode !== Mode.Framework) {
+      return;
+    }
+
+    const modelsAsDataLanguage = getModelsAsDataLanguage(this.language);
+    const modelGeneration = modelsAsDataLanguage.modelGeneration;
+    const autoRun = modelGeneration?.autoRun;
+
+    if (modelGeneration === undefined || autoRun === undefined) {
+      return;
+    }
+
+    await withProgress(
+      async (progress) => {
+        progress({
+          step: 0,
+          maxStep: 4000,
+          message: "Generating models",
+        });
+
+        const parseResults =
+          autoRun.parseResults ?? modelGeneration.parseResults;
+
+        try {
+          await runGenerateQueries({
+            queryConstraints: modelGeneration.queryConstraints(mode),
+            filterQueries: modelGeneration.filterQueries,
+            parseResults: (queryPath, results) =>
+              parseResults(
+                queryPath,
+                results,
+                modelsAsDataLanguage,
+                this.app.logger,
+              ),
+            onResults: async (modeledMethods) => {
+              this.addModeledMethodsFromArray(modeledMethods);
+            },
+            cliServer: this.cliServer,
+            queryRunner: this.queryRunner,
+            queryStorageDir: this.queryStorageDir,
+            databaseItem: this.databaseItem,
+            progress,
+            token: this.cancellationTokenSource.token,
+          });
+        } catch (e: unknown) {
+          void showAndLogExceptionWithTelemetry(
+            this.app.logger,
+            this.app.telemetry,
+            redactableError(
+              asError(e),
+            )`Failed to auto-run generating models: ${getErrorMessage(e)}`,
+          );
+        }
+      },
+      {
+        cancellable: false,
+        location: ProgressLocation.Window,
+        title: "Generating models",
+      },
+    );
+  }
+
   private async generateModeledMethodsFromLlm(
     packageName: string,
     methodSignatures: string[],