Add generating of flow model to data extension editor

koesie10 · koesie10 · commit 7f65122adb82 · 2023-04-04T13:54:35.000+02:00
This adds the automatic generation of sources/sinks/summary flows to the data extension editor using the flow model queries. This is based on the Python script available in the CodeQL repo. See: https://github.com/github/codeql/blob/main/java/ql/src/utils/modelgenerator/GenerateFlowModel.py
diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-module.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-module.ts
@@ -5,10 +5,12 @@ import { CodeQLCliServer } from "../cli";
 import { QueryRunner } from "../queryRunner";
 import { DatabaseManager } from "../local-databases";
 import { extLogger } from "../common";
+import { App } from "../common/app";
 
 export class DataExtensionsEditorModule {
   public constructor(
     private readonly ctx: ExtensionContext,
+    private readonly app: App,
     private readonly databaseManager: DatabaseManager,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
@@ -26,6 +28,8 @@ export class DataExtensionsEditorModule {
 
         const view = new DataExtensionsEditorView(
           this.ctx,
+          this.app,
+          this.databaseManager,
           this.cliServer,
           this.queryRunner,
           this.queryStorageDir,
diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts
@@ -18,16 +18,22 @@ import { file } from "tmp-promise";
 import { readFile, writeFile } from "fs-extra";
 import { dump, load } from "js-yaml";
 import { getOnDiskWorkspaceFolders } from "../helpers";
-import { DatabaseItem } from "../local-databases";
+import { DatabaseItem, DatabaseManager } from "../local-databases";
 import { CodeQLCliServer } from "../cli";
-import { assertNever } from "../pure/helpers-pure";
+import { assertNever, getErrorMessage } from "../pure/helpers-pure";
+import { generateFlowModel } from "./generate-flow-model";
+import { ModeledMethod } from "./interface";
+import { promptImportGithubDatabase } from "../databaseFetcher";
+import { App } from "../common/app";
 
 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,
@@ -69,6 +75,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
         await this.saveYaml(msg.yaml);
         await this.loadExternalApiUsages();
 
+        break;
+      case "generateExternalApi":
+        await this.generateExternalApi();
+
         break;
       default:
         assertNever(msg);
@@ -149,6 +159,84 @@ export class DataExtensionsEditorView extends AbstractWebview<
     await this.clearProgress();
   }
 
+  protected async generateExternalApi(): Promise<void> {
+    const tokenSource = new CancellationTokenSource();
+
+    const selectedDatabase = this.databaseManager.currentDatabaseItem;
+
+    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;
+    }
+
+    await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
+
+    const workspaceFolder = workspace.workspaceFolders?.find(
+      (folder) => folder.name === "ql",
+    );
+    if (!workspaceFolder) {
+      void extLogger.log("No workspace folder 'ql' found");
+
+      return;
+    }
+
+    await this.showProgress({
+      step: 0,
+      maxStep: 4000,
+      message: "Generating external API",
+    });
+
+    try {
+      await generateFlowModel(
+        this.cliServer,
+        this.queryRunner,
+        this.queryStorageDir,
+        workspaceFolder.uri.fsPath,
+        database,
+        async (results) => {
+          const modeledMethodsByName: Record<string, ModeledMethod> = {};
+
+          for (const result of results) {
+            modeledMethodsByName[result[0]] = result[1];
+          }
+
+          await this.postMessage({
+            t: "addModeledMethods",
+            modeledMethods: modeledMethodsByName,
+          });
+        },
+        (update) => this.showProgress(update),
+        tokenSource.token,
+      );
+    } catch (e: unknown) {
+      void extLogger.log(`Error: ${getErrorMessage(e)}`);
+    }
+
+    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);
 
diff --git a/extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts b/extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts
@@ -0,0 +1,146 @@
+import { CancellationToken } from "vscode";
+import { DatabaseItem } from "../local-databases";
+import { ModeledMethod, ModeledMethodType } from "./interface";
+import { join } from "path";
+import { QueryRunner } from "../queryRunner";
+import { CodeQLCliServer } from "../cli";
+import { extLogger, TeeLogger } from "../common";
+import { definitions } from "./yaml";
+import { ProgressCallback } from "../progress";
+import { getOnDiskWorkspaceFolders } from "../helpers";
+
+class FlowModelGenerator {
+  constructor(
+    private readonly cli: CodeQLCliServer,
+    private readonly queryRunner: QueryRunner,
+    private readonly queryStorageDir: string,
+    private readonly qlDir: string,
+    private readonly databaseItem: DatabaseItem,
+    private readonly progress: ProgressCallback,
+    private readonly token: CancellationToken,
+  ) {}
+
+  async getAddsTo(
+    type: Exclude<ModeledMethodType, "none">,
+    queryName: string,
+    queryStep: number,
+  ): Promise<Array<[string, ModeledMethod]> | undefined> {
+    const definition = definitions[type];
+
+    const query = join(
+      this.qlDir,
+      this.databaseItem.language,
+      "ql/src/utils/modelgenerator",
+      queryName,
+    );
+
+    const queryRun = this.queryRunner.createQueryRun(
+      this.databaseItem.databaseUri.fsPath,
+      { queryPath: query, quickEvalPosition: undefined },
+      false,
+      getOnDiskWorkspaceFolders(),
+      undefined,
+      this.queryStorageDir,
+      undefined,
+      undefined,
+    );
+
+    const queryResult = await queryRun.evaluate(
+      ({ step, message }) =>
+        this.progress({
+          message: `Generating ${type} model: ${message}`,
+          step: queryStep * 1000 + step,
+          maxStep: 4000,
+        }),
+      this.token,
+      new TeeLogger(this.queryRunner.logger, queryRun.outputDir.logPath),
+    );
+
+    const bqrsPath = queryResult.outputDir.bqrsPath;
+
+    const bqrsInfo = await this.cli.bqrsInfo(bqrsPath);
+    if (bqrsInfo["result-sets"].length !== 1) {
+      void extLogger.log(
+        `Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
+      );
+      return undefined;
+    }
+
+    const resultSet = bqrsInfo["result-sets"][0];
+
+    const decodedResults = await this.cli.bqrsDecode(bqrsPath, resultSet.name);
+
+    const results = decodedResults.tuples;
+
+    return results
+      .map((result): [string, ModeledMethod] | undefined => {
+        const row = result[0] as string;
+
+        return definition.readModeledMethod(row.split(";"));
+      })
+      .filter(
+        (result): result is [string, ModeledMethod] => result !== undefined,
+      );
+  }
+
+  async run(
+    onResults: (
+      results: Array<[string, ModeledMethod]>,
+    ) => void | Promise<void>,
+  ) {
+    const summaryResults = await this.getAddsTo(
+      "summary",
+      "CaptureSummaryModels.ql",
+      0,
+    );
+    if (summaryResults) {
+      await onResults(summaryResults);
+    }
+
+    const sinkResults = await this.getAddsTo("sink", "CaptureSinkModels.ql", 1);
+    if (sinkResults) {
+      await onResults(sinkResults);
+    }
+
+    const sourceResults = await this.getAddsTo(
+      "source",
+      "CaptureSourceModels.ql",
+      2,
+    );
+    if (sourceResults) {
+      await onResults(sourceResults);
+    }
+
+    const neutralResults = await this.getAddsTo(
+      "neutral",
+      "CaptureNeutralModels.ql",
+      3,
+    );
+    if (neutralResults) {
+      await onResults(neutralResults);
+    }
+  }
+}
+
+export async function generateFlowModel(
+  cli: CodeQLCliServer,
+  queryRunner: QueryRunner,
+  queryStorageDir: string,
+  qlDir: string,
+  databaseItem: DatabaseItem,
+  onResults: (results: Array<[string, ModeledMethod]>) => void | Promise<void>,
+  progress: ProgressCallback,
+  token: CancellationToken,
+) {
+  const generator = new FlowModelGenerator(
+    cli,
+    queryRunner,
+    queryStorageDir,
+    qlDir,
+    databaseItem,
+    progress,
+    token,
+  );
+
+  return generator.run(onResults);
+}
diff --git a/extensions/ql-vscode/src/data-extensions-editor/yaml.ts b/extensions/ql-vscode/src/data-extensions-editor/yaml.ts
@@ -19,7 +19,7 @@ function readRowToMethod(row: any[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }
 
-const definitions: Record<
+export const definitions: Record<
   Exclude<ModeledMethodType, "none">,
   DataExtensionDefinition
 > = {
diff --git a/extensions/ql-vscode/src/extension.ts b/extensions/ql-vscode/src/extension.ts
@@ -868,6 +868,7 @@ async function activateWithInstalledDistribution(
   await ensureDir(dataExtensionsEditorQueryStorageDir);
   const dataExtensionsEditorModule = new DataExtensionsEditorModule(
     ctx,
+    app,
     dbm,
     cliServer,
     qs,
diff --git a/extensions/ql-vscode/src/pure/interface-types.ts b/extensions/ql-vscode/src/pure/interface-types.ts
@@ -15,6 +15,7 @@ import {
 import { RepositoriesFilterSortStateWithIds } from "./variant-analysis-filter-sort";
 import { ErrorLike } from "./errors";
 import { DataFlowPaths } from "../variant-analysis/shared/data-flow-paths";
+import { ModeledMethod } from "../data-extensions-editor/interface";
 
 /**
  * This module contains types and code that are shared between
@@ -497,16 +498,27 @@ export interface SetExistingYamlDataMessage {
   data: any;
 }
 
+export interface AddModeledMethodsMessage {
+  t: "addModeledMethods";
+  modeledMethods: Record<string, ModeledMethod>;
+}
+
 export interface ApplyDataExtensionYamlMessage {
   t: "applyDataExtensionYaml";
   yaml: string;
 }
 
+export interface GenerateExternalApiMessage {
+  t: "generateExternalApi";
+}
+
 export type ToDataExtensionsEditorMessage =
   | SetExternalApiResultsMessage
   | ShowProgressMessage
-  | SetExistingYamlDataMessage;
+  | SetExistingYamlDataMessage
+  | AddModeledMethodsMessage;
 
 export type FromDataExtensionsEditorMessage =
   | ViewLoadedMsg
-  | ApplyDataExtensionYamlMessage;
+  | ApplyDataExtensionYamlMessage
+  | GenerateExternalApiMessage;
diff --git a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
@@ -74,6 +74,20 @@ export function DataExtensionsEditor(): JSX.Element {
               };
             });
 
+            break;
+          case "addModeledMethods":
+            setModeledMethods((oldModeledMethods) => {
+              const filteredOldModeledMethods = Object.fromEntries(
+                Object.entries(oldModeledMethods).filter(
+                  ([, value]) => value.type !== "none",
+                ),
+              );
+
+              return {
+                ...msg.modeledMethods,
+                ...filteredOldModeledMethods,
+              };
+            });
             break;
           default:
             assertNever(msg);
@@ -168,6 +182,12 @@ export function DataExtensionsEditor(): JSX.Element {
     });
   }, [methods, modeledMethods]);
 
+  const onGenerateClick = useCallback(() => {
+    vscode.postMessage({
+      t: "generateExternalApi",
+    });
+  }, []);
+
   return (
     <DataExtensionsEditorContainer>
       {progress.maxStep > 0 && (
@@ -189,6 +209,12 @@ export function DataExtensionsEditor(): JSX.Element {
           <div>
             <h3>External API modelling</h3>
             <VSCodeButton onClick={onApplyClick}>Apply</VSCodeButton>
+            &nbsp;
+            <VSCodeButton onClick={onGenerateClick}>
+              Download and generate
+            </VSCodeButton>
+            <br />
+            <br />
             <VSCodeDataGrid>
               <VSCodeDataGridRow rowType="header">
                 <VSCodeDataGridCell cellType="columnheader" gridColumn={1}>

Original file line number	Diff line number	Diff line change
`@@ -19,7 +19,7 @@ function readRowToMethod(row: any[]): string {`
`19`	`19`	return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
`20`	`20`	`}`
`21`	`21`
`22`		`-const definitions: Record<`
	`22`	`+export const definitions: Record<`
`23`	`23`	`Exclude<ModeledMethodType, "none">,`
`24`	`24`	`DataExtensionDefinition`
`25`	`25`	`> = {`