Show external API calls in data extensions editor

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

@@ -0,0 +1,54 @@
+import { DecodedBqrsChunk } from "../pure/bqrs-cli-types";
+import { Call, ExternalApiUsage } from "./external-api-usage";
+
+export function decodeBqrsToExternalApiUsages(
+  chunk: DecodedBqrsChunk,
+): ExternalApiUsage[] {
+  const methodsByApiName = new Map<string, ExternalApiUsage>();
+
+  chunk?.tuples.forEach((tuple) => {
+    const signature = tuple[0] as string;
+    const supported = tuple[1] as boolean;
+    const usage = tuple[2] as Call;
+
+    const [packageWithType, methodDeclaration] = signature.split("#");
+
+    const packageName = packageWithType.substring(
+      0,
+      packageWithType.lastIndexOf("."),
+    );
+    const typeName = packageWithType.substring(
+      packageWithType.lastIndexOf(".") + 1,
+    );
+
+    const methodName = methodDeclaration.substring(
+      0,
+      methodDeclaration.indexOf("("),
+    );
+    const methodParameters = methodDeclaration.substring(
+      methodDeclaration.indexOf("("),
+    );
+
+    if (!methodsByApiName.has(signature)) {
+      methodsByApiName.set(signature, {
+        signature,
+        packageName,
+        typeName,
+        methodName,
+        methodParameters,
+        supported,
+        usages: [],
+      });
+    }
+
+    const method = methodsByApiName.get(signature)!;
+    method.usages.push(usage);
+  });
+
+  const externalApiUsages = Array.from(methodsByApiName.values());
+  externalApiUsages.sort((a, b) => {
+    // Sort by number of usages descending
+    return b.usages.length - a.usages.length;
+  });
+  return externalApiUsages;
+}

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

@@ -1,16 +1,70 @@
 import { ExtensionContext } from "vscode";
 import { DataExtensionsEditorView } from "./data-extensions-editor-view";
 import { DataExtensionsEditorCommands } from "../common/commands";
+import { CodeQLCliServer } from "../cli";
+import { QueryRunner } from "../queryRunner";
+import { DatabaseManager } from "../local-databases";
+import { extLogger } from "../common";
+import { ensureDir } from "fs-extra";
+import { join } from "path";
 
 export class DataExtensionsEditorModule {
-  public constructor(private readonly ctx: ExtensionContext) {}
+  private readonly queryStorageDir: string;
+
+  private constructor(
+    private readonly ctx: ExtensionContext,
+    private readonly databaseManager: DatabaseManager,
+    private readonly cliServer: CodeQLCliServer,
+    private readonly queryRunner: QueryRunner,
+    baseQueryStorageDir: string,
+  ) {
+    this.queryStorageDir = join(
+      baseQueryStorageDir,
+      "data-extensions-editor-results",
+    );
+  }
+
+  public static async initialize(
+    ctx: ExtensionContext,
+    databaseManager: DatabaseManager,
+    cliServer: CodeQLCliServer,
+    queryRunner: QueryRunner,
+    queryStorageDir: string,
+  ): Promise<DataExtensionsEditorModule> {
+    const dataExtensionsEditorModule = new DataExtensionsEditorModule(
+      ctx,
+      databaseManager,
+      cliServer,
+      queryRunner,
+      queryStorageDir,
+    );
+
+    await dataExtensionsEditorModule.initialize();
+    return dataExtensionsEditorModule;
+  }
 
   public getCommands(): DataExtensionsEditorCommands {
     return {
       "codeQL.openDataExtensionsEditor": async () => {
-        const view = new DataExtensionsEditorView(this.ctx);
+        const db = this.databaseManager.currentDatabaseItem;
+        if (!db) {
+          void extLogger.log("No database selected");
+          return;
+        }
+
+        const view = new DataExtensionsEditorView(
+          this.ctx,
+          this.cliServer,
+          this.queryRunner,
+          this.queryStorageDir,
+          db,
+        );
         await view.openView();
       },
     };
   }
+
+  private async initialize(): Promise<void> {
+    await ensureDir(this.queryStorageDir);
+  }
 }

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

@@ -1,15 +1,37 @@
-import { ExtensionContext, ViewColumn } from "vscode";
+import { CancellationTokenSource, ExtensionContext, ViewColumn } from "vscode";
 import { AbstractWebview, WebviewPanelConfig } from "../abstract-webview";
 import {
   FromDataExtensionsEditorMessage,
   ToDataExtensionsEditorMessage,
 } from "../pure/interface-types";
+import { ProgressUpdate } from "../progress";
+import { extLogger, TeeLogger } from "../common";
+import { CoreCompletedQuery, QueryRunner } from "../queryRunner";
+import { qlpackOfDatabase } from "../contextual/queryResolver";
+import { file } from "tmp-promise";
+import { writeFile } from "fs-extra";
+import { dump } from "js-yaml";
+import {
+  getOnDiskWorkspaceFolders,
+  showAndLogExceptionWithTelemetry,
+} from "../helpers";
+import { DatabaseItem } from "../local-databases";
+import { CodeQLCliServer } from "../cli";
+import { decodeBqrsToExternalApiUsages } from "./bqrs";
+import { redactableError } from "../pure/errors";
+import { asError, getErrorMessage } from "../pure/helpers-pure";
 
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
   FromDataExtensionsEditorMessage
 > {
-  public constructor(ctx: ExtensionContext) {
+  public constructor(
+    ctx: ExtensionContext,
+    private readonly cliServer: CodeQLCliServer,
+    private readonly queryRunner: QueryRunner,
+    private readonly queryStorageDir: string,
+    private readonly databaseItem: DatabaseItem,
+  ) {
     super(ctx);
   }
 
@@ -49,5 +71,154 @@ export class DataExtensionsEditorView extends AbstractWebview<
 
   protected async onWebViewLoaded() {
     super.onWebViewLoaded();
+
+    await this.loadExternalApiUsages();
+  }
+
+  protected async loadExternalApiUsages(): Promise<void> {
+    try {
+      const queryResult = await this.runQuery();
+      if (!queryResult) {
+        await this.clearProgress();
+        return;
+      }
+
+      await this.showProgress({
+        message: "Loading results",
+        step: 1100,
+        maxStep: 1500,
+      });
+
+      const bqrsPath = queryResult.outputDir.bqrsPath;
+
+      const bqrsChunk = await this.getResults(bqrsPath);
+      if (!bqrsChunk) {
+        await this.clearProgress();
+        return;
+      }
+
+      await this.showProgress({
+        message: "Finalizing results",
+        step: 1450,
+        maxStep: 1500,
+      });
+
+      const externalApiUsages = decodeBqrsToExternalApiUsages(bqrsChunk);
+
+      await this.postMessage({
+        t: "setExternalApiUsages",
+        externalApiUsages,
+      });
+
+      await this.clearProgress();
+    } catch (err) {
+      void showAndLogExceptionWithTelemetry(
+        redactableError(
+          asError(err),
+        )`Failed to load external APi usages: ${getErrorMessage(err)}`,
+      );
+    }
+  }
+
+  private async runQuery(): Promise<CoreCompletedQuery | undefined> {
+    const qlpacks = await qlpackOfDatabase(this.cliServer, this.databaseItem);
+
+    const packsToSearch = [qlpacks.dbschemePack];
+    if (qlpacks.queryPack) {
+      packsToSearch.push(qlpacks.queryPack);
+    }
+
+    const suiteFile = (
+      await file({
+        postfix: ".qls",
+      })
+    ).path;
+    const suiteYaml = [];
+    for (const qlpack of packsToSearch) {
+      suiteYaml.push({
+        from: qlpack,
+        queries: ".",
+        include: {
+          id: `${this.databaseItem.language}/telemetry/fetch-external-apis`,
+        },
+      });
+    }
+    await writeFile(suiteFile, dump(suiteYaml), "utf8");
+
+    const queries = await this.cliServer.resolveQueriesInSuite(
+      suiteFile,
+      getOnDiskWorkspaceFolders(),
+    );
+
+    if (queries.length !== 1) {
+      void extLogger.log(`Expected exactly one query, got ${queries.length}`);
+      return;
+    }
+
+    const query = queries[0];
+
+    const tokenSource = new CancellationTokenSource();
+
+    const queryRun = this.queryRunner.createQueryRun(
+      this.databaseItem.databaseUri.fsPath,
+      { queryPath: query, quickEvalPosition: undefined },
+      false,
+      getOnDiskWorkspaceFolders(),
+      undefined,
+      this.queryStorageDir,
+      undefined,
+      undefined,
+    );
+
+    return queryRun.evaluate(
+      (update) => this.showProgress(update, 1500),
+      tokenSource.token,
+      new TeeLogger(this.queryRunner.logger, queryRun.outputDir.logPath),
+    );
+  }
+
+  private async getResults(bqrsPath: string) {
+    const bqrsInfo = await this.cliServer.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];
+
+    await this.showProgress({
+      message: "Decoding results",
+      step: 1200,
+      maxStep: 1500,
+    });
+
+    return this.cliServer.bqrsDecode(bqrsPath, resultSet.name);
+  }
+
+  /*
+   * Progress in this class is a bit weird. Most of the progress is based on running the query.
+   * Query progress is always between 0 and 1000. However, we still have some steps that need
+   * to be done after the query has finished. Therefore, the maximum step is 1500. This captures
+   * 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.
+   */
+  private async showProgress(update: ProgressUpdate, maxStep?: number) {
+    await this.postMessage({
+      t: "showProgress",
+      step: update.step,
+      maxStep: maxStep ?? update.maxStep,
+      message: update.message,
+    });
+  }
+
+  private async clearProgress() {
+    await this.showProgress({
+      step: 0,
+      maxStep: 0,
+      message: "",
+    });
   }
 }

extensions/ql-vscode/src/data-extensions-editor/external-api-usage.ts

@@ -0,0 +1,19 @@
+import { ResolvableLocationValue } from "../pure/bqrs-cli-types";
+
+export type Call = {
+  label: string;
+  url: ResolvableLocationValue;
+};
+
+export type ExternalApiUsage = {
+  /**
+   * Contains the full method signature, e.g. `org.sql2o.Connection#createQuery(String)`
+   */
+  signature: string;
+  packageName: string;
+  typeName: string;
+  methodName: string;
+  methodParameters: string;
+  supported: boolean;
+  usages: Call[];
+};

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

@@ -0,0 +1,13 @@
+export type ModeledMethodType =
+  | "none"
+  | "source"
+  | "sink"
+  | "summary"
+  | "neutral";
+
+export type ModeledMethod = {
+  type: ModeledMethodType;
+  input: string;
+  output: string;
+  kind: string;
+};

extensions/ql-vscode/src/extension.ts

@@ -865,7 +865,14 @@ async function activateWithInstalledDistribution(
   );
   ctx.subscriptions.push(localQueries);
 
-  const dataExtensionsEditorModule = new DataExtensionsEditorModule(ctx);
+  const dataExtensionsEditorModule =
+    await DataExtensionsEditorModule.initialize(
+      ctx,
+      dbm,
+      cliServer,
+      qs,
+      tmpDir.name,
+    );
 
   void extLogger.log("Initializing QLTest interface.");
   const testExplorerExtension = extensions.getExtension<TestHub>(