From 398f16c4cbf149c878fd23a3b5018bbc3b93e4b9 Mon Sep 17 00:00:00 2001
From: Anders Starcke Henriksen <starcke@github.com>
Date: Wed, 23 Aug 2023 13:00:03 +0200
Subject: [PATCH] Change fetch queries to use shared methods.

---
 .../data-extensions-editor-view.ts            | 29 +------
 .../external-api-usage-queries.ts             | 79 ++++++++++---------
 2 files changed, 44 insertions(+), 64 deletions(-)

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
index 631780e7a..8f2dd3d84 100644
--- 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
@@ -26,12 +26,8 @@ import { runFlowModelQueries } from "./flow-model-queries";
 import { promptImportGithubDatabase } from "../databases/database-fetcher";
 import { App } from "../common/app";
 import { showResolvableLocation } from "../databases/local-databases/locations";
-import { decodeBqrsToExternalApiUsages } from "./bqrs";
 import { redactableError } from "../common/errors";
-import {
-  readQueryResults,
-  runExternalApiQueries,
-} from "./external-api-usage-queries";
+import { runExternalApiQueries } from "./external-api-usage-queries";
 import { ExternalApiUsage, Usage } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 import { ExtensionPack } from "./shared/extension-pack";
@@ -319,28 +315,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
           if (!queryResult) {
             return;
           }
-
-          progress({
-            message: "Decoding results",
-            step: 1100,
-            maxStep: 1500,
-          });
-
-          const bqrsChunk = await readQueryResults({
-            cliServer: this.cliServer,
-            bqrsPath: queryResult.outputDir.bqrsPath,
-          });
-          if (!bqrsChunk) {
-            return;
-          }
-
-          progress({
-            message: "Finalizing results",
-            step: 1450,
-            maxStep: 1500,
-          });
-
-          this.externalApiUsages = decodeBqrsToExternalApiUsages(bqrsChunk);
+          this.externalApiUsages = queryResult;
 
           await this.postMessage({
             t: "setExternalApiUsages",
diff --git a/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-queries.ts b/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-queries.ts
index be1f2b0e7..90da8b111 100644
--- a/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-queries.ts
+++ b/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-queries.ts
@@ -1,12 +1,11 @@
-import { CoreCompletedQuery, QueryRunner } from "../query-server";
+import { QueryRunner } from "../query-server";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { extLogger } from "../common/logging/vscode";
-import { showAndLogExceptionWithTelemetry, TeeLogger } from "../common/logging";
+import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { CancellationToken } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { DatabaseItem } from "../databases/local-databases";
 import { ProgressCallback } from "../common/vscode/progress";
-import { QueryResultType } from "../query-server/new-messages";
 import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { join } from "path";
@@ -14,11 +13,14 @@ import { Mode } from "./shared/mode";
 import { writeFile } from "fs-extra";
 import { QueryLanguage } from "../common/query-language";
 import { fetchExternalApiQueries } from "./queries";
+import { ExternalApiUsage } from "./external-api-usage";
+import { runQuery } from "../local-queries/run-query";
+import { decodeBqrsToExternalApiUsages } from "./bqrs";
 
 type RunQueryOptions = {
-  cliServer: Pick<CodeQLCliServer, "resolveQlpacks">;
-  queryRunner: Pick<QueryRunner, "createQueryRun" | "logger">;
-  databaseItem: Pick<DatabaseItem, "contents" | "databaseUri" | "language">;
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  databaseItem: DatabaseItem;
   queryStorageDir: string;
   queryDir: string;
 
@@ -67,7 +69,7 @@ export async function runExternalApiQueries(
     progress,
     token,
   }: RunQueryOptions,
-): Promise<CoreCompletedQuery | undefined> {
+): Promise<ExternalApiUsage[] | undefined> {
   // The below code is temporary to allow for rapid prototyping of the queries. Once the queries are stabilized, we will
   // move these queries into the `github/codeql` repository and use them like any other contextual (e.g. AST) queries.
   // This is intentionally not pretty code, as it will be removed soon.
@@ -79,44 +81,47 @@ export async function runExternalApiQueries(
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
-  const queryFile = join(
-    queryDir,
-    `FetchExternalApis${mode.charAt(0).toUpperCase() + mode.slice(1)}Mode.ql`,
-  );
+  const queryPath = join(queryDir, queryNameFromMode(mode));
 
-  const queryRun = queryRunner.createQueryRun(
-    databaseItem.databaseUri.fsPath,
-    {
-      queryPath: queryFile,
-      quickEvalPosition: undefined,
-      quickEvalCountOnly: false,
-    },
-    false,
-    getOnDiskWorkspaceFolders(),
-    extensionPacks,
+  // Run the actual query
+  const completedQuery = await runQuery({
+    cliServer,
+    queryRunner,
+    databaseItem,
+    queryPath,
     queryStorageDir,
-    undefined,
-    undefined,
-  );
-
-  const completedQuery = await queryRun.evaluate(
+    additionalPacks,
+    extensionPacks,
     progress,
     token,
-    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
-  );
+  });
 
-  if (completedQuery.resultType !== QueryResultType.SUCCESS) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`External API usage query failed: ${
-        completedQuery.message ?? "No message"
-      }`,
-    );
+  if (!completedQuery) {
     return;
   }
 
-  return completedQuery;
+  // Read the results and covert to internal representation
+  progress({
+    message: "Decoding results",
+    step: 1100,
+    maxStep: 1500,
+  });
+
+  const bqrsChunk = await readQueryResults({
+    cliServer,
+    bqrsPath: completedQuery.outputDir.bqrsPath,
+  });
+  if (!bqrsChunk) {
+    return;
+  }
+
+  progress({
+    message: "Finalizing results",
+    step: 1450,
+    maxStep: 1500,
+  });
+
+  return decodeBqrsToExternalApiUsages(bqrsChunk);
 }
 
 type GetResultsOptions = {