Merge remote-tracking branch 'origin/main' into dbartol/debug-adapter

2023-04-11 16:06:25 +00:00
parent 1cbfd0159e b0940e6da9
commit e8f39fea83
20 changed files with 1029 additions and 132 deletions
--- a/extensions/ql-vscode/.babelrc.json
+++ b/extensions/ql-vscode/.babelrc.json
@@ -0,0 +1,16 @@
+{
+  "sourceType": "unambiguous",
+  "presets": [
+    [
+      "@babel/preset-env",
+      {
+        "targets": {
+          "chrome": 100
+        }
+      }
+    ],
+    "@babel/preset-typescript",
+    "@babel/preset-react"
+  ],
+  "plugins": []
+}
--- a/extensions/ql-vscode/.storybook/main.ts
+++ b/extensions/ql-vscode/.storybook/main.ts
@@ -12,6 +12,9 @@ const config: StorybookConfig = {
  core: {
    builder: "@storybook/builder-webpack5",
  },
+  features: {
+    babelModeV7: true,
+  },
 };

 module.exports = config;
--- a/extensions/ql-vscode/CHANGELOG.md
+++ b/extensions/ql-vscode/CHANGELOG.md
@@ -2,6 +2,7 @@

 ## [UNRELEASED]

+- Fix bug that was causing code flows to not get updated when switching between results. [#2288](https://github.com/github/vscode-codeql/pull/2288)
 - Restart the CodeQL language server whenever the _CodeQL: Restart Query Server_ command is invoked. This avoids bugs where the CLI version changes to support new language features, but the language server is not updated. [#2238](https://github.com/github/vscode-codeql/pull/2238)

 ## 1.8.1 - 23 March 2023
--- a/extensions/ql-vscode/src/contextual/queryResolver.ts
+++ b/extensions/ql-vscode/src/contextual/queryResolver.ts
@@ -21,8 +21,8 @@ import { redactableError } from "../pure/errors";
 import { QLPACK_FILENAMES } from "../pure/ql";

 export async function qlpackOfDatabase(
-  cli: CodeQLCliServer,
-  db: DatabaseItem,
+  cli: Pick<CodeQLCliServer, "resolveQlpacks">,
+  db: Pick<DatabaseItem, "contents">,
 ): Promise<QlPacksForLanguage> {
  if (db.contents === undefined) {
    throw new Error("Database is invalid and cannot infer QLPack.");
--- 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
@@ -1,18 +1,20 @@
 import { ExtensionContext } from "vscode";
 import { DataExtensionsEditorView } from "./data-extensions-editor-view";
 import { DataExtensionsEditorCommands } from "../common/commands";
-import { CodeQLCliServer } from "../cli";
+import { CliVersionConstraint, 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";
+import { App } from "../common/app";
+import { showAndLogErrorMessage } from "../helpers";

 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,6 +28,7 @@ export class DataExtensionsEditorModule {

  public static async initialize(
    ctx: ExtensionContext,
+    app: App,
    databaseManager: DatabaseManager,
    cliServer: CodeQLCliServer,
    queryRunner: QueryRunner,
@@ -33,6 +36,7 @@ export class DataExtensionsEditorModule {
  ): Promise<DataExtensionsEditorModule> {
    const dataExtensionsEditorModule = new DataExtensionsEditorModule(
      ctx,
+      app,
      databaseManager,
      cliServer,
      queryRunner,
@@ -48,12 +52,21 @@ export class DataExtensionsEditorModule {
      "codeQL.openDataExtensionsEditor": async () => {
        const db = this.databaseManager.currentDatabaseItem;
        if (!db) {
-          void extLogger.log("No database selected");
+          void showAndLogErrorMessage("No database selected");
+          return;
+        }
+
+        if (!(await this.cliServer.cliConstraints.supportsQlpacksKind())) {
+          void showAndLogErrorMessage(
+            `This feature requires CodeQL CLI version ${CliVersionConstraint.CLI_VERSION_WITH_QLPACKS_KIND.format()} or later.`,
+          );
          return;
        }

        const view = new DataExtensionsEditorView(
          this.ctx,
+          this.app,
+          this.databaseManager,
          this.cliServer,
          this.queryRunner,
          this.queryStorageDir,
--- 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
@@ -5,6 +5,7 @@ import {
  ViewColumn,
  window,
  workspace,
+  WorkspaceFolder,
 } from "vscode";
 import { AbstractWebview, WebviewPanelConfig } from "../abstract-webview";
 import {
@@ -12,34 +13,50 @@ import {
  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 { readFile, writeFile } from "fs-extra";
-import { dump as dumpYaml, load as loadYaml } from "js-yaml";
+import { QueryRunner } from "../queryRunner";
 import {
-  getOnDiskWorkspaceFolders,
  showAndLogExceptionWithTelemetry,
  showAndLogWarningMessage,
 } from "../helpers";
-import { DatabaseItem } from "../local-databases";
+import { extLogger } from "../common";
+import { readFile, writeFile } from "fs-extra";
+import { load as loadYaml } from "js-yaml";
+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";
 import { redactableError } from "../pure/errors";
+import { readQueryResults, runQuery } from "./external-api-usage-query";
 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 +105,10 @@ export class DataExtensionsEditorView extends AbstractWebview<
        );
        await this.loadExternalApiUsages();

+        break;
+      case "generateExternalApi":
+        await this.generateModeledMethods();
+
        break;
      default:
        assertNever(msg);
@@ -160,8 +181,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}`);
@@ -169,22 +190,36 @@ export class DataExtensionsEditorView extends AbstractWebview<
  }

  protected async loadExternalApiUsages(): Promise<void> {
+    const cancellationTokenSource = new CancellationTokenSource();
+
    try {
-      const queryResult = await this.runQuery();
+      const queryResult = await runQuery({
+        cliServer: this.cliServer,
+        queryRunner: this.queryRunner,
+        databaseItem: this.databaseItem,
+        queryStorageDir: this.queryStorageDir,
+        logger: extLogger,
+        progress: (progressUpdate: ProgressUpdate) => {
+          void this.showProgress(progressUpdate, 1500);
+        },
+        token: cancellationTokenSource.token,
+      });
      if (!queryResult) {
        await this.clearProgress();
        return;
      }

      await this.showProgress({
-        message: "Loading results",
+        message: "Decoding results",
        step: 1100,
        maxStep: 1500,
      });

-      const bqrsPath = queryResult.outputDir.bqrsPath;
-
-      const bqrsChunk = await this.getResults(bqrsPath);
+      const bqrsChunk = await readQueryResults({
+        cliServer: this.cliServer,
+        bqrsPath: queryResult.outputDir.bqrsPath,
+        logger: extLogger,
+      });
      if (!bqrsChunk) {
        await this.clearProgress();
        return;
@@ -213,81 +248,90 @@ export class DataExtensionsEditorView extends AbstractWebview<
    }
  }

-  private async runQuery(): Promise<CoreCompletedQuery | undefined> {
-    const qlpacks = await qlpackOfDatabase(this.cliServer, this.databaseItem);
+  protected async generateModeledMethods(): Promise<void> {
+    const tokenSource = new CancellationTokenSource();

-    const packsToSearch = [qlpacks.dbschemePack];
-    if (qlpacks.queryPack) {
-      packsToSearch.push(qlpacks.queryPack);
-    }
+    const selectedDatabase = this.databaseManager.currentDatabaseItem;

-    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, dumpYaml(suiteYaml), "utf8");
-
-    const queries = await this.cliServer.resolveQueriesInSuite(
-      suiteFile,
-      getOnDiskWorkspaceFolders(),
+    // 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");

-    if (queries.length !== 1) {
-      void extLogger.log(`Expected exactly one query, got ${queries.length}`);
      return;
    }

-    const query = queries[0];
+    // 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 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),
-    );
+    const workspaceFolder = getQlSubmoduleFolder();
+    if (!workspaceFolder) {
+      return;
    }

-  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,
+      step: 0,
+      maxStep: 4000,
+      message: "Generating modeled methods for library",
    });

-    return this.cliServer.bqrsDecode(bqrsPath, resultSet.name);
+    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();
  }

  /*
@@ -297,6 +341,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 +367,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;
    }

--- a/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts
+++ b/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts
@@ -0,0 +1,114 @@
+import { CoreCompletedQuery, QueryRunner } from "../queryRunner";
+import { qlpackOfDatabase } from "../contextual/queryResolver";
+import { file } from "tmp-promise";
+import { writeFile } from "fs-extra";
+import { dump as dumpYaml } from "js-yaml";
+import { getOnDiskWorkspaceFolders } from "../helpers";
+import { Logger, TeeLogger } from "../common";
+import { CancellationToken } from "vscode";
+import { CodeQLCliServer } from "../cli";
+import { DatabaseItem } from "../local-databases";
+import { ProgressCallback } from "../progress";
+
+export type RunQueryOptions = {
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks" | "resolveQueriesInSuite">;
+  queryRunner: Pick<QueryRunner, "createQueryRun" | "logger">;
+  databaseItem: Pick<DatabaseItem, "contents" | "databaseUri" | "language">;
+  queryStorageDir: string;
+  logger: Logger;
+
+  progress: ProgressCallback;
+  token: CancellationToken;
+};
+
+export async function runQuery({
+  cliServer,
+  queryRunner,
+  databaseItem,
+  queryStorageDir,
+  logger,
+  progress,
+  token,
+}: RunQueryOptions): Promise<CoreCompletedQuery | undefined> {
+  const qlpacks = await qlpackOfDatabase(cliServer, 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: `${databaseItem.language}/telemetry/fetch-external-apis`,
+      },
+    });
+  }
+  await writeFile(suiteFile, dumpYaml(suiteYaml), "utf8");
+
+  const additionalPacks = getOnDiskWorkspaceFolders();
+  const extensionPacks = Object.keys(
+    await cliServer.resolveQlpacks(additionalPacks, true),
+  );
+
+  const queries = await cliServer.resolveQueriesInSuite(
+    suiteFile,
+    getOnDiskWorkspaceFolders(),
+  );
+
+  if (queries.length !== 1) {
+    void logger.log(`Expected exactly one query, got ${queries.length}`);
+    return;
+  }
+
+  const query = queries[0];
+
+  const queryRun = queryRunner.createQueryRun(
+    databaseItem.databaseUri.fsPath,
+    { queryPath: query, quickEvalPosition: undefined },
+    false,
+    getOnDiskWorkspaceFolders(),
+    extensionPacks,
+    queryStorageDir,
+    undefined,
+    undefined,
+  );
+
+  return queryRun.evaluate(
+    progress,
+    token,
+    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
+  );
+}
+
+export type GetResultsOptions = {
+  cliServer: Pick<CodeQLCliServer, "bqrsInfo" | "bqrsDecode">;
+  bqrsPath: string;
+  logger: Logger;
+};
+
+export async function readQueryResults({
+  cliServer,
+  bqrsPath,
+  logger,
+}: GetResultsOptions) {
+  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
+  if (bqrsInfo["result-sets"].length !== 1) {
+    void logger.log(
+      `Expected exactly one result set, got ${bqrsInfo["result-sets"].length}`,
+    );
+    return undefined;
+  }
+
+  const resultSet = bqrsInfo["result-sets"][0];
+
+  return cliServer.bqrsDecode(bqrsPath, resultSet.name);
+}
--- 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,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);
+  }
+}
--- a/extensions/ql-vscode/src/data-extensions-editor/modeled-method.ts
+++ b/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;
+};
--- a/extensions/ql-vscode/src/data-extensions-editor/yaml.ts
+++ b/extensions/ql-vscode/src/data-extensions-editor/yaml.ts
@@ -1,27 +1,31 @@
 import { ExternalApiUsage } from "./external-api-usage";
-import { ModeledMethod, ModeledMethodType } from "./modeled-method";
+import {
+  ModeledMethod,
+  ModeledMethodType,
+  ModeledMethodWithSignature,
+} from "./modeled-method";

 type ExternalApiUsageByType = {
  externalApiUsage: ExternalApiUsage;
  modeledMethod: ModeledMethod;
 };

-type DataExtensionDefinition = {
-  extensible: string;
+type ExtensiblePredicateDefinition = {
+  extensiblePredicate: string;
  generateMethodDefinition: (method: ExternalApiUsageByType) => any[];
-  readModeledMethod: (row: any[]) => [string, ModeledMethod] | undefined;
+  readModeledMethod: (row: any[]) => ModeledMethodWithSignature;
 };

 function readRowToMethod(row: any[]): string {
  return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }

-const definitions: Record<
+export const extensiblePredicateDefinitions: Record<
  Exclude<ModeledMethodType, "none">,
-  DataExtensionDefinition
+  ExtensiblePredicateDefinition
 > = {
  source: {
-    extensible: "sourceModel",
+    extensiblePredicate: "sourceModel",
    // extensible predicate sourceModel(
    //   string package, string type, boolean subtypes, string name, string signature, string ext,
    //   string output, string kind, string provenance
@@ -37,18 +41,18 @@ const definitions: Record<
      method.modeledMethod.kind,
      "manual",
    ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
        type: "source",
        input: "",
        output: row[6],
        kind: row[7],
      },
-    ],
+    }),
  },
  sink: {
-    extensible: "sinkModel",
+    extensiblePredicate: "sinkModel",
    // extensible predicate sinkModel(
    //   string package, string type, boolean subtypes, string name, string signature, string ext,
    //   string input, string kind, string provenance
@@ -64,18 +68,18 @@ const definitions: Record<
      method.modeledMethod.kind,
      "manual",
    ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
        type: "sink",
        input: row[6],
        output: "",
        kind: row[7],
      },
-    ],
+    }),
  },
  summary: {
-    extensible: "summaryModel",
+    extensiblePredicate: "summaryModel",
    // extensible predicate summaryModel(
    //   string package, string type, boolean subtypes, string name, string signature, string ext,
    //   string input, string output, string kind, string provenance
@@ -92,18 +96,18 @@ const definitions: Record<
      method.modeledMethod.kind,
      "manual",
    ],
-    readModeledMethod: (row) => [
-      readRowToMethod(row),
-      {
+    readModeledMethod: (row) => ({
+      signature: readRowToMethod(row),
+      modeledMethod: {
        type: "summary",
        input: row[6],
        output: row[7],
        kind: row[8],
      },
-    ],
+    }),
  },
  neutral: {
-    extensible: "neutralModel",
+    extensiblePredicate: "neutralModel",
    // extensible predicate neutralModel(
    //   string package, string type, string name, string signature, string provenance
    // );
@@ -114,21 +118,21 @@ const definitions: Record<
      method.externalApiUsage.methodParameters,
      "manual",
    ],
-    readModeledMethod: (row) => [
-      `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
-      {
+    readModeledMethod: (row) => ({
+      signature: `${row[0]}.${row[1]}#${row[2]}${row[3]}`,
+      modeledMethod: {
        type: "neutral",
        input: "",
        output: "",
        kind: "",
      },
-    ],
+    }),
  },
 };

 function createDataProperty(
  methods: ExternalApiUsageByType[],
-  definition: DataExtensionDefinition,
+  definition: ExtensiblePredicateDefinition,
 ) {
  if (methods.length === 0) {
    return " []";
@@ -169,10 +173,10 @@ export function createDataExtensionYaml(
    }
  }

-  const extensions = Object.entries(definitions).map(
+  const extensions = Object.entries(extensiblePredicateDefinitions).map(
    ([type, definition]) => `  - addsTo:
      pack: codeql/java-all
-      extensible: ${definition.extensible}
+      extensible: ${definition.extensiblePredicate}
    data:${createDataProperty(
      methodsByType[type as Exclude<ModeledMethodType, "none">],
      definition,
@@ -214,8 +218,8 @@ export function loadDataExtensionYaml(
      continue;
    }

-    const definition = Object.values(definitions).find(
-      (definition) => definition.extensible === extensible,
+    const definition = Object.values(extensiblePredicateDefinitions).find(
+      (definition) => definition.extensiblePredicate === extensible,
    );
    if (!definition) {
      continue;
@@ -227,9 +231,9 @@ export function loadDataExtensionYaml(
        continue;
      }

-      const [apiInfo, modeledMethod] = result;
+      const { signature, modeledMethod } = result;

-      modeledMethods[apiInfo] = modeledMethod;
+      modeledMethods[signature] = modeledMethod;
    }
  }

--- a/extensions/ql-vscode/src/distribution.ts
+++ b/extensions/ql-vscode/src/distribution.ts
@@ -215,6 +215,9 @@ export class DistributionManager implements DistributionProvider {
    minSecondsSinceLastUpdateCheck: number,
  ): Promise<DistributionUpdateCheckResult> {
    const distribution = await this.getDistributionWithoutVersionCheck();
+    if (distribution === undefined) {
+      minSecondsSinceLastUpdateCheck = 0;
+    }
    const extensionManagedCodeQlPath =
      await this.extensionSpecificDistributionManager.getCodeQlPathWithoutVersionCheck();
    if (distribution?.codeQlPath !== extensionManagedCodeQlPath) {
--- a/extensions/ql-vscode/src/extension.ts
+++ b/extensions/ql-vscode/src/extension.ts
@@ -884,6 +884,7 @@ async function activateWithInstalledDistribution(
  const dataExtensionsEditorModule =
    await DataExtensionsEditorModule.initialize(
      ctx,
+      app,
      dbm,
      cliServer,
      qs,
--- a/extensions/ql-vscode/src/helpers.ts
+++ b/extensions/ql-vscode/src/helpers.ts
@@ -478,7 +478,7 @@ function findStandardQueryPack(
 }

 export async function getQlPackForDbscheme(
-  cliServer: CodeQLCliServer,
+  cliServer: Pick<CodeQLCliServer, "resolveQlpacks">,
  dbschemePath: string,
 ): Promise<QlPacksForLanguage> {
  const qlpacks = await cliServer.resolveQlpacks(getOnDiskWorkspaceFolders());
--- a/extensions/ql-vscode/src/pure/interface-types.ts
+++ b/extensions/ql-vscode/src/pure/interface-types.ts
@@ -493,6 +493,19 @@ export interface ShowProgressMessage {
  message: string;
 }

+export interface AddModeledMethodsMessage {
+  t: "addModeledMethods";
+  modeledMethods: Record<string, ModeledMethod>;
+
+  /**
+   * If true, then any existing modeled methods set to "none" will be
+   * overwritten by the new modeled methods. Otherwise, the "none" modeled
+   * methods will not be overwritten, even if the new modeled methods
+   * contain a better model.
+   */
+  overrideNone?: boolean;
+}
+
 export interface JumpToUsageMessage {
  t: "jumpToUsage";
  location: ResolvableLocationValue;
@@ -509,12 +522,17 @@ export interface SaveModeledMethods {
  modeledMethods: Record<string, ModeledMethod>;
 }

+export interface GenerateExternalApiMessage {
+  t: "generateExternalApi";
+}
+
 export type ToDataExtensionsEditorMessage =
  | SetExternalApiUsagesMessage
  | ShowProgressMessage
-  | SetExistingModeledMethods;
+  | AddModeledMethodsMessage;

 export type FromDataExtensionsEditorMessage =
  | ViewLoadedMsg
  | JumpToUsageMessage
-  | SaveModeledMethods;
+  | SaveModeledMethods
+  | GenerateExternalApiMessage;
--- a/extensions/ql-vscode/src/stories/data-extensions-editor/DataExtensionsEditor.stories.tsx
+++ b/extensions/ql-vscode/src/stories/data-extensions-editor/DataExtensionsEditor.stories.tsx
@@ -0,0 +1,223 @@
+import * as React from "react";
+
+import { ComponentMeta, ComponentStory } from "@storybook/react";
+
+import { DataExtensionsEditor as DataExtensionsEditorComponent } from "../../view/data-extensions-editor/DataExtensionsEditor";
+
+export default {
+  title: "Data Extensions Editor/Data Extensions Editor",
+  component: DataExtensionsEditorComponent,
+} as ComponentMeta<typeof DataExtensionsEditorComponent>;
+
+const Template: ComponentStory<typeof DataExtensionsEditorComponent> = (
+  args,
+) => <DataExtensionsEditorComponent {...args} />;
+
+export const DataExtensionsEditor = Template.bind({});
+DataExtensionsEditor.args = {
+  initialExternalApiUsages: [
+    {
+      signature: "org.sql2o.Connection#createQuery(String)",
+      packageName: "org.sql2o",
+      typeName: "Connection",
+      methodName: "createQuery",
+      methodParameters: "(String)",
+      supported: true,
+      usages: [
+        {
+          label: "createQuery(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 15,
+            startColumn: 13,
+            endLine: 15,
+            endColumn: 56,
+          },
+        },
+        {
+          label: "createQuery(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 26,
+            startColumn: 13,
+            endLine: 26,
+            endColumn: 39,
+          },
+        },
+      ],
+    },
+    {
+      signature: "org.sql2o.Query#executeScalar(Class)",
+      packageName: "org.sql2o",
+      typeName: "Query",
+      methodName: "executeScalar",
+      methodParameters: "(Class)",
+      supported: true,
+      usages: [
+        {
+          label: "executeScalar(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 15,
+            startColumn: 13,
+            endLine: 15,
+            endColumn: 85,
+          },
+        },
+        {
+          label: "executeScalar(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 26,
+            startColumn: 13,
+            endLine: 26,
+            endColumn: 68,
+          },
+        },
+      ],
+    },
+    {
+      signature: "org.sql2o.Sql2o#open()",
+      packageName: "org.sql2o",
+      typeName: "Sql2o",
+      methodName: "open",
+      methodParameters: "()",
+      supported: false,
+      usages: [
+        {
+          label: "open(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 14,
+            startColumn: 24,
+            endLine: 14,
+            endColumn: 35,
+          },
+        },
+        {
+          label: "open(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 25,
+            startColumn: 24,
+            endLine: 25,
+            endColumn: 35,
+          },
+        },
+      ],
+    },
+    {
+      signature: "java.io.PrintStream#println(String)",
+      packageName: "java.io",
+      typeName: "PrintStream",
+      methodName: "println",
+      methodParameters: "(String)",
+      supported: true,
+      usages: [
+        {
+          label: "println(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 29,
+            startColumn: 9,
+            endLine: 29,
+            endColumn: 49,
+          },
+        },
+      ],
+    },
+    {
+      signature:
+        "org.springframework.boot.SpringApplication#run(Class,String[])",
+      packageName: "org.springframework.boot",
+      typeName: "SpringApplication",
+      methodName: "run",
+      methodParameters: "(Class,String[])",
+      supported: false,
+      usages: [
+        {
+          label: "run(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/Sql2oExampleApplication.java",
+            startLine: 9,
+            startColumn: 9,
+            endLine: 9,
+            endColumn: 66,
+          },
+        },
+      ],
+    },
+    {
+      signature: "org.sql2o.Sql2o#Sql2o(String,String,String)",
+      packageName: "org.sql2o",
+      typeName: "Sql2o",
+      methodName: "Sql2o",
+      methodParameters: "(String,String,String)",
+      supported: false,
+      usages: [
+        {
+          label: "new Sql2o(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 10,
+            startColumn: 33,
+            endLine: 10,
+            endColumn: 88,
+          },
+        },
+      ],
+    },
+    {
+      signature: "org.sql2o.Sql2o#Sql2o(String)",
+      packageName: "org.sql2o",
+      typeName: "Sql2o",
+      methodName: "Sql2o",
+      methodParameters: "(String)",
+      supported: false,
+      usages: [
+        {
+          label: "new Sql2o(...)",
+          url: {
+            uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+            startLine: 23,
+            startColumn: 23,
+            endLine: 23,
+            endColumn: 36,
+          },
+        },
+      ],
+    },
+  ],
+  initialModeledMethods: {
+    "org.sql2o.Sql2o#Sql2o(String)": {
+      type: "sink",
+      input: "Argument[0]",
+      output: "",
+      kind: "jndi-injection",
+    },
+    "org.sql2o.Connection#createQuery(String)": {
+      type: "summary",
+      input: "Argument[-1]",
+      output: "ReturnValue",
+      kind: "taint",
+    },
+    "org.sql2o.Sql2o#open()": {
+      type: "summary",
+      input: "Argument[-1]",
+      output: "ReturnValue",
+      kind: "taint",
+    },
+    "org.sql2o.Query#executeScalar(Class)": {
+      type: "neutral",
+      input: "",
+      output: "",
+      kind: "",
+    },
+    "org.sql2o.Sql2o#Sql2o(String,String,String)": {
+      type: "neutral",
+      input: "",
+      output: "",
+      kind: "",
+    },
+  },
+};
--- a/extensions/ql-vscode/src/stories/data-extensions-editor/MethodRow.stories.tsx
+++ b/extensions/ql-vscode/src/stories/data-extensions-editor/MethodRow.stories.tsx
@@ -0,0 +1,54 @@
+import * as React from "react";
+
+import { ComponentMeta, ComponentStory } from "@storybook/react";
+
+import { MethodRow as MethodRowComponent } from "../../view/data-extensions-editor/MethodRow";
+
+export default {
+  title: "Data Extensions Editor/Method Row",
+  component: MethodRowComponent,
+} as ComponentMeta<typeof MethodRowComponent>;
+
+const Template: ComponentStory<typeof MethodRowComponent> = (args) => (
+  <MethodRowComponent {...args} />
+);
+
+export const MethodRow = Template.bind({});
+MethodRow.args = {
+  externalApiUsage: {
+    signature: "org.sql2o.Sql2o#open()",
+    packageName: "org.sql2o",
+    typeName: "Sql2o",
+    methodName: "open",
+    methodParameters: "()",
+    supported: true,
+    usages: [
+      {
+        label: "open(...)",
+        url: {
+          uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+          startLine: 14,
+          startColumn: 24,
+          endLine: 14,
+          endColumn: 35,
+        },
+      },
+      {
+        label: "open(...)",
+        url: {
+          uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+          startLine: 25,
+          startColumn: 24,
+          endLine: 25,
+          endColumn: 35,
+        },
+      },
+    ],
+  },
+  modeledMethod: {
+    type: "summary",
+    input: "Argument[-1]",
+    output: "ReturnValue",
+    kind: "taint",
+  },
+};
--- a/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
+++ b/extensions/ql-vscode/src/view/data-extensions-editor/DataExtensionsEditor.tsx
@@ -33,13 +33,21 @@ const ProgressBar = styled.div<ProgressBarProps>`
  background-color: var(--vscode-progressBar-background);
 `;

-export function DataExtensionsEditor(): JSX.Element {
+type Props = {
+  initialExternalApiUsages?: ExternalApiUsage[];
+  initialModeledMethods?: Record<string, ModeledMethod>;
+};
+
+export function DataExtensionsEditor({
+  initialExternalApiUsages = [],
+  initialModeledMethods = {},
+}: Props): JSX.Element {
  const [externalApiUsages, setExternalApiUsages] = useState<
    ExternalApiUsage[]
-  >([]);
+  >(initialExternalApiUsages);
  const [modeledMethods, setModeledMethods] = useState<
    Record<string, ModeledMethod>
-  >({});
+  >(initialModeledMethods);
  const [progress, setProgress] = useState<Omit<ShowProgressMessage, "t">>({
    step: 0,
    maxStep: 0,
@@ -57,14 +65,21 @@ export function DataExtensionsEditor(): JSX.Element {
          case "showProgress":
            setProgress(msg);
            break;
-          case "setExistingModeledMethods":
+          case "addModeledMethods":
            setModeledMethods((oldModeledMethods) => {
+              const filteredOldModeledMethods = msg.overrideNone
+                ? Object.fromEntries(
+                    Object.entries(oldModeledMethods).filter(
+                      ([, value]) => value.type !== "none",
+                    ),
+                  )
+                : oldModeledMethods;
+
              return {
-                ...msg.existingModeledMethods,
-                ...oldModeledMethods,
+                ...msg.modeledMethods,
+                ...filteredOldModeledMethods,
              };
            });
-
            break;
          default:
            assertNever(msg);
@@ -107,6 +122,12 @@ export function DataExtensionsEditor(): JSX.Element {
    });
  }, [externalApiUsages, modeledMethods]);

+  const onGenerateClick = useCallback(() => {
+    vscode.postMessage({
+      t: "generateExternalApi",
+    });
+  }, []);
+
  return (
    <DataExtensionsEditorContainer>
      {progress.maxStep > 0 && (
@@ -128,6 +149,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}>
--- a/extensions/ql-vscode/src/view/data-flow-paths/DataFlowPaths.tsx
+++ b/extensions/ql-vscode/src/view/data-flow-paths/DataFlowPaths.tsx
@@ -41,6 +41,11 @@ export const DataFlowPaths = ({

  const { codeFlows, ruleDescription, message, severity } = dataFlowPaths;

+  React.useEffect(() => {
+    // Make sure to update the selected code flow if the data flow paths change
+    setSelectedCodeFlow(dataFlowPaths.codeFlows[0]);
+  }, [dataFlowPaths]);
+
  return (
    <>
      <VerticalSpace size={2} />
--- a/extensions/ql-vscode/src/view/results/result-tables.tsx
+++ b/extensions/ql-vscode/src/view/results/result-tables.tsx
@@ -79,7 +79,7 @@ function renderResultCountString(resultSet: ResultSet): JSX.Element {
 }

 function getInterpretedTableName(interpretation: Interpretation): string {
-  return interpretation?.data.t === "GraphInterpretationData"
+  return interpretation.data.t === "GraphInterpretationData"
    ? GRAPH_TABLE_NAME
    : ALERTS_TABLE_NAME;
 }
@@ -101,7 +101,7 @@ function getResultSets(
 ): ResultSet[] {
  const resultSets: ResultSet[] =
    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore 2783
+    // @ts-ignore 2783 Avoid compilation error for overwriting the t property
    rawResultSets.map((rs) => ({ t: "RawResultSet", ...rs }));

  if (interpretation !== undefined) {
--- a/extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts
+++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts
@@ -0,0 +1,222 @@
+import {
+  readQueryResults,
+  runQuery,
+} from "../../../../src/data-extensions-editor/external-api-usage-query";
+import { createMockLogger } from "../../../__mocks__/loggerMock";
+import type { Uri } from "vscode";
+import { DatabaseKind } from "../../../../src/local-databases";
+import * as queryResolver from "../../../../src/contextual/queryResolver";
+import { file } from "tmp-promise";
+import { QueryResultType } from "../../../../src/pure/new-messages";
+import { readFile } from "fs-extra";
+import { load } from "js-yaml";
+
+function createMockUri(path = "/a/b/c/foo"): Uri {
+  return {
+    scheme: "file",
+    authority: "",
+    path,
+    query: "",
+    fragment: "",
+    fsPath: path,
+    with: jest.fn(),
+    toJSON: jest.fn(),
+  };
+}
+
+describe("runQuery", () => {
+  it("runs the query", async () => {
+    jest.spyOn(queryResolver, "qlpackOfDatabase").mockResolvedValue({
+      dbschemePack: "codeql/java-all",
+      dbschemePackIsLibraryPack: false,
+      queryPack: "codeql/java-queries",
+    });
+
+    const logPath = (await file()).path;
+
+    const options = {
+      cliServer: {
+        resolveQlpacks: jest.fn().mockResolvedValue({
+          "my/java-extensions": "/a/b/c/",
+        }),
+        resolveQueriesInSuite: jest
+          .fn()
+          .mockResolvedValue([
+            "/home/github/codeql/java/ql/src/Telemetry/FetchExternalAPIs.ql",
+          ]),
+      },
+      queryRunner: {
+        createQueryRun: jest.fn().mockReturnValue({
+          evaluate: jest.fn().mockResolvedValue({
+            resultType: QueryResultType.SUCCESS,
+          }),
+          outputDir: {
+            logPath,
+          },
+        }),
+        logger: createMockLogger(),
+      },
+      logger: createMockLogger(),
+      databaseItem: {
+        databaseUri: createMockUri("/a/b/c/src.zip"),
+        contents: {
+          kind: DatabaseKind.Database,
+          name: "foo",
+          datasetUri: createMockUri(),
+        },
+        language: "java",
+      },
+      queryStorageDir: "/tmp/queries",
+      progress: jest.fn(),
+      token: {
+        isCancellationRequested: false,
+        onCancellationRequested: jest.fn(),
+      },
+    };
+    const result = await runQuery(options);
+
+    expect(result?.resultType).toEqual(QueryResultType.SUCCESS);
+
+    expect(options.cliServer.resolveQueriesInSuite).toHaveBeenCalledWith(
+      expect.anything(),
+      [],
+    );
+    const suiteFile = options.cliServer.resolveQueriesInSuite.mock.calls[0][0];
+    const suiteFileContents = await readFile(suiteFile, "utf8");
+    const suiteYaml = load(suiteFileContents);
+    expect(suiteYaml).toEqual([
+      {
+        from: "codeql/java-all",
+        queries: ".",
+        include: {
+          id: "java/telemetry/fetch-external-apis",
+        },
+      },
+      {
+        from: "codeql/java-queries",
+        queries: ".",
+        include: {
+          id: "java/telemetry/fetch-external-apis",
+        },
+      },
+    ]);
+
+    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledTimes(1);
+    expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith([], true);
+    expect(options.queryRunner.createQueryRun).toHaveBeenCalledWith(
+      "/a/b/c/src.zip",
+      {
+        queryPath:
+          "/home/github/codeql/java/ql/src/Telemetry/FetchExternalAPIs.ql",
+        quickEvalPosition: undefined,
+      },
+      false,
+      [],
+      ["my/java-extensions"],
+      "/tmp/queries",
+      undefined,
+      undefined,
+    );
+  });
+});
+
+describe("readQueryResults", () => {
+  const options = {
+    cliServer: {
+      bqrsInfo: jest.fn(),
+      bqrsDecode: jest.fn(),
+    },
+    bqrsPath: "/tmp/results.bqrs",
+    logger: createMockLogger(),
+  };
+
+  it("returns undefined when there are no results", async () => {
+    options.cliServer.bqrsInfo.mockResolvedValue({
+      "result-sets": [],
+    });
+
+    expect(await readQueryResults(options)).toBeUndefined();
+    expect(options.logger.log).toHaveBeenCalledWith(
+      expect.stringMatching(/Expected exactly one result set/),
+    );
+  });
+
+  it("returns undefined when there are multiple result sets", async () => {
+    options.cliServer.bqrsInfo.mockResolvedValue({
+      "result-sets": [
+        {
+          name: "#select",
+          rows: 10,
+          columns: [
+            { name: "apiName", kind: "s" },
+            { name: "supported", kind: "b" },
+            { name: "usage", kind: "e" },
+          ],
+        },
+        {
+          name: "#select2",
+          rows: 10,
+          columns: [
+            { name: "apiName", kind: "s" },
+            { name: "supported", kind: "b" },
+            { name: "usage", kind: "e" },
+          ],
+        },
+      ],
+    });
+
+    expect(await readQueryResults(options)).toBeUndefined();
+    expect(options.logger.log).toHaveBeenCalledWith(
+      expect.stringMatching(/Expected exactly one result set/),
+    );
+  });
+
+  it("gets the result set", async () => {
+    options.cliServer.bqrsInfo.mockResolvedValue({
+      "result-sets": [
+        {
+          name: "#select",
+          rows: 10,
+          columns: [
+            { name: "apiName", kind: "s" },
+            { name: "supported", kind: "b" },
+            { name: "usage", kind: "e" },
+          ],
+        },
+      ],
+      "compatible-query-kinds": ["Table", "Tree", "Graph"],
+    });
+    const decodedResultSet = {
+      columns: [
+        { name: "apiName", kind: "String" },
+        { name: "supported", kind: "Boolean" },
+        { name: "usage", kind: "Entity" },
+      ],
+      tuples: [
+        [
+          "java.io.PrintStream#println(String)",
+          true,
+          {
+            label: "println(...)",
+            url: {
+              uri: "file:/home/runner/work/sql2o-example/sql2o-example/src/main/java/org/example/HelloController.java",
+              startLine: 29,
+              startColumn: 9,
+              endLine: 29,
+              endColumn: 49,
+            },
+          },
+        ],
+      ],
+    };
+    options.cliServer.bqrsDecode.mockResolvedValue(decodedResultSet);
+
+    const result = await readQueryResults(options);
+    expect(result).toEqual(decodedResultSet);
+    expect(options.cliServer.bqrsInfo).toHaveBeenCalledWith(options.bqrsPath);
+    expect(options.cliServer.bqrsDecode).toHaveBeenCalledWith(
+      options.bqrsPath,
+      "#select",
+    );
+  });
+});