Merge pull request #2267 from github/koesie10/data-extension-editor-generate-flow-model

Add generating of flow model to data extension editor

Author: koesie10

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

@@ -7,12 +7,14 @@ import { DatabaseManager } from "../local-databases";
 import { extLogger } from "../common";
 import { ensureDir } from "fs-extra";
 import { join } from "path";
+import { App } from "../common/app";
 
 export class DataExtensionsEditorModule {
   private readonly queryStorageDir: string;
 
   private constructor(
     private readonly ctx: ExtensionContext,
+    private readonly app: App,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -26,13 +28,15 @@ export class DataExtensionsEditorModule {
 
   public static async initialize(
     ctx: ExtensionContext,
+    app: App,
     databaseManager: DatabaseManager,
     cliServer: CodeQLCliServer,
     queryRunner: QueryRunner,
     queryStorageDir: string,
   ): Promise<DataExtensionsEditorModule> {
     const dataExtensionsEditorModule = new DataExtensionsEditorModule(
       ctx,
+      app,
       databaseManager,
       cliServer,
       queryRunner,
@@ -54,6 +58,8 @@ export class DataExtensionsEditorModule {
 
         const view = new DataExtensionsEditorView(
           this.ctx,
+          this.app,
+          this.databaseManager,
           this.cliServer,
           this.queryRunner,
           this.queryStorageDir,

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

@@ -5,6 +5,7 @@ import {
   ViewColumn,
   window,
   workspace,
+  WorkspaceFolder,
 } from "vscode";
 import { AbstractWebview, WebviewPanelConfig } from "../abstract-webview";
 import {
@@ -23,9 +24,12 @@ import {
   showAndLogExceptionWithTelemetry,
   showAndLogWarningMessage,
 } from "../helpers";
-import { DatabaseItem } from "../local-databases";
+import { DatabaseItem, DatabaseManager } from "../local-databases";
 import { CodeQLCliServer } from "../cli";
 import { asError, assertNever, getErrorMessage } from "../pure/helpers-pure";
+import { generateFlowModel } from "./generate-flow-model";
+import { promptImportGithubDatabase } from "../databaseFetcher";
+import { App } from "../common/app";
 import { ResolvableLocationValue } from "../pure/bqrs-cli-types";
 import { showResolvableLocation } from "../interface-utils";
 import { decodeBqrsToExternalApiUsages } from "./bqrs";
@@ -34,12 +38,27 @@ import { createDataExtensionYaml, loadDataExtensionYaml } from "./yaml";
 import { ExternalApiUsage } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 
+function getQlSubmoduleFolder(): WorkspaceFolder | undefined {
+  const workspaceFolder = workspace.workspaceFolders?.find(
+    (folder) => folder.name === "ql",
+  );
+  if (!workspaceFolder) {
+    void extLogger.log("No workspace folder 'ql' found");
+
+    return;
+  }
+
+  return workspaceFolder;
+}
+
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
   FromDataExtensionsEditorMessage
 > {
   public constructor(
     ctx: ExtensionContext,
+    private readonly app: App,
+    private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
     private readonly queryStorageDir: string,
@@ -88,6 +107,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
         );
         await this.loadExternalApiUsages();
 
+        break;
+      case "generateExternalApi":
+        await this.generateModeledMethods();
+
         break;
       default:
         assertNever(msg);
@@ -160,8 +183,8 @@ export class DataExtensionsEditorView extends AbstractWebview<
       }
 
       await this.postMessage({
-        t: "setExistingModeledMethods",
-        existingModeledMethods,
+        t: "addModeledMethods",
+        modeledMethods: existingModeledMethods,
       });
     } catch (e: unknown) {
       void extLogger.log(`Unable to read data extension YAML: ${e}`);
@@ -213,6 +236,92 @@ export class DataExtensionsEditorView extends AbstractWebview<
     }
   }
 
+  protected async generateModeledMethods(): Promise<void> {
+    const tokenSource = new CancellationTokenSource();
+
+    const selectedDatabase = this.databaseManager.currentDatabaseItem;
+
+    // The external API methods are in the library source code, so we need to ask
+    // the user to import the library database. We need to have the database
+    // imported to the query server, so we need to register it to our workspace.
+    const database = await promptImportGithubDatabase(
+      this.app.commands,
+      this.databaseManager,
+      this.app.workspaceStoragePath ?? this.app.globalStoragePath,
+      this.app.credentials,
+      (update) => this.showProgress(update),
+      tokenSource.token,
+      this.cliServer,
+    );
+    if (!database) {
+      await this.clearProgress();
+      void extLogger.log("No database chosen");
+
+      return;
+    }
+
+    // The library database was set as the current database by importing it,
+    // but we need to set it back to the originally selected database.
+    await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
+
+    const workspaceFolder = getQlSubmoduleFolder();
+    if (!workspaceFolder) {
+      return;
+    }
+
+    await this.showProgress({
+      step: 0,
+      maxStep: 4000,
+      message: "Generating modeled methods for library",
+    });
+
+    try {
+      await generateFlowModel({
+        cliServer: this.cliServer,
+        queryRunner: this.queryRunner,
+        queryStorageDir: this.queryStorageDir,
+        qlDir: workspaceFolder.uri.fsPath,
+        databaseItem: database,
+        onResults: async (results) => {
+          const modeledMethodsByName: Record<string, ModeledMethod> = {};
+
+          for (const result of results) {
+            modeledMethodsByName[result.signature] = result.modeledMethod;
+          }
+
+          await this.postMessage({
+            t: "addModeledMethods",
+            modeledMethods: modeledMethodsByName,
+            overrideNone: true,
+          });
+        },
+        progress: (update) => this.showProgress(update),
+        token: tokenSource.token,
+      });
+    } catch (e: unknown) {
+      void showAndLogExceptionWithTelemetry(
+        redactableError(
+          asError(e),
+        )`Failed to generate flow model: ${getErrorMessage(e)}`,
+      );
+    }
+
+    // After the flow model has been generated, we can remove the temporary database
+    // which we used for generating the flow model.
+    await this.databaseManager.removeDatabaseItem(
+      () =>
+        this.showProgress({
+          step: 3900,
+          maxStep: 4000,
+          message: "Removing temporary database",
+        }),
+      tokenSource.token,
+      database,
+    );
+
+    await this.clearProgress();
+  }
+
   private async runQuery(): Promise<CoreCompletedQuery | undefined> {
     const qlpacks = await qlpackOfDatabase(this.cliServer, this.databaseItem);
 
@@ -297,6 +406,13 @@ export class DataExtensionsEditorView extends AbstractWebview<
    * that there's 1000 steps of the query progress since that takes the most time, and then
    * an additional 500 steps for the rest of the work. The progress doesn't need to be 100%
    * accurate, so this is just a rough estimate.
+   *
+   * For generating the modeled methods for an external library, the max step is 4000. This is
+   * based on the following steps:
+   * - 1000 for the summary model
+   * - 1000 for the sink model
+   * - 1000 for the source model
+   * - 1000 for the neutral model
    */
   private async showProgress(update: ProgressUpdate, maxStep?: number) {
     await this.postMessage({
@@ -316,12 +432,8 @@ export class DataExtensionsEditorView extends AbstractWebview<
   }
 
   private calculateModelFilename(): string | undefined {
-    const workspaceFolder = workspace.workspaceFolders?.find(
-      (folder) => folder.name === "ql",
-    );
+    const workspaceFolder = getQlSubmoduleFolder();
     if (!workspaceFolder) {
-      void extLogger.log("No workspace folder 'ql' found");
-
       return;
     }
 

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

@@ -0,0 +1,141 @@
+import { CancellationToken } from "vscode";
+import { DatabaseItem } from "../local-databases";
+import { join } from "path";
+import { QueryRunner } from "../queryRunner";
+import { CodeQLCliServer } from "../cli";
+import { TeeLogger } from "../common";
+import { extensiblePredicateDefinitions } from "./yaml";
+import { ProgressCallback } from "../progress";
+import { getOnDiskWorkspaceFolders } from "../helpers";
+import {
+  ModeledMethodType,
+  ModeledMethodWithSignature,
+} from "./modeled-method";
+
+type FlowModelOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  queryStorageDir: string;
+  qlDir: string;
+  databaseItem: DatabaseItem;
+  progress: ProgressCallback;
+  token: CancellationToken;
+  onResults: (results: ModeledMethodWithSignature[]) => void | Promise<void>;
+};
+
+async function getModeledMethodsFromFlow(
+  type: Exclude<ModeledMethodType, "none">,
+  queryName: string,
+  queryStep: number,
+  {
+    cliServer,
+    queryRunner,
+    queryStorageDir,
+    qlDir,
+    databaseItem,
+    progress,
+    token,
+  }: Omit<FlowModelOptions, "onResults">,
+): Promise<ModeledMethodWithSignature[]> {
+  const definition = extensiblePredicateDefinitions[type];
+
+  const query = join(
+    qlDir,
+    databaseItem.language,
+    "ql/src/utils/modelgenerator",
+    queryName,
+  );
+
+  const queryRun = queryRunner.createQueryRun(
+    databaseItem.databaseUri.fsPath,
+    { queryPath: query, quickEvalPosition: undefined },
+    false,
+    getOnDiskWorkspaceFolders(),
+    undefined,
+    queryStorageDir,
+    undefined,
+    undefined,
+  );
+
+  const queryResult = await queryRun.evaluate(
+    ({ step, message }) =>
+      progress({
+        message: `Generating ${type} model: ${message}`,
+        step: queryStep * 1000 + step,
+        maxStep: 4000,
+      }),
+    token,
+    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
+  );
+
+  const bqrsPath = queryResult.outputDir.bqrsPath;
+
+  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
+  if (bqrsInfo["result-sets"].length !== 1) {
+    throw new Error(
+      `Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
+    );
+  }
+
+  const resultSet = bqrsInfo["result-sets"][0];
+
+  const decodedResults = await cliServer.bqrsDecode(bqrsPath, resultSet.name);
+
+  const results = decodedResults.tuples;
+
+  return (
+    results
+      // This is just a sanity check. The query should only return strings.
+      .filter((result) => typeof result[0] === "string")
+      .map((result) => {
+        const row = result[0] as string;
+
+        return definition.readModeledMethod(row.split(";"));
+      })
+  );
+}
+
+export async function generateFlowModel({
+  onResults,
+  ...options
+}: FlowModelOptions) {
+  const summaryResults = await getModeledMethodsFromFlow(
+    "summary",
+    "CaptureSummaryModels.ql",
+    0,
+    options,
+  );
+  if (summaryResults) {
+    await onResults(summaryResults);
+  }
+
+  const sinkResults = await getModeledMethodsFromFlow(
+    "sink",
+    "CaptureSinkModels.ql",
+    1,
+    options,
+  );
+  if (sinkResults) {
+    await onResults(sinkResults);
+  }
+
+  const sourceResults = await getModeledMethodsFromFlow(
+    "source",
+    "CaptureSourceModels.ql",
+    2,
+    options,
+  );
+  if (sourceResults) {
+    await onResults(sourceResults);
+  }
+
+  const neutralResults = await getModeledMethodsFromFlow(
+    "neutral",
+    "CaptureNeutralModels.ql",
+    3,
+    options,
+  );
+  if (neutralResults) {
+    await onResults(neutralResults);
+  }
+}

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

@@ -11,3 +11,8 @@ export type ModeledMethod = {
   output: string;
   kind: string;
 };
+
+export type ModeledMethodWithSignature = {
+  signature: string;
+  modeledMethod: ModeledMethod;
+};