Starcke/fetch queries out of view

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

@@ -26,9 +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, runQuery } 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";
@@ -304,7 +303,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
       async (progress) => {
         try {
           const cancellationTokenSource = new CancellationTokenSource();
-          const queryResult = await runQuery(this.mode, {
+          const queryResult = await runExternalApiQueries(this.mode, {
             cliServer: this.cliServer,
             queryRunner: this.queryRunner,
             databaseItem: this.databaseItem,
@@ -316,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",

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

@@ -1,24 +1,26 @@
-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";
 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;
 
@@ -56,7 +58,7 @@ export async function prepareExternalApiQuery(
   return true;
 }
 
-export async function runQuery(
+export async function runExternalApiQueries(
   mode: Mode,
   {
     cliServer,
@@ -67,7 +69,7 @@ export async function runQuery(
     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 runQuery(
     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;
+  }
+
+  // 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;
   }
 
-  return completedQuery;
+  progress({
+    message: "Finalizing results",
+    step: 1450,
+    maxStep: 1500,
+  });
+
+  return decodeBqrsToExternalApiUsages(bqrsChunk);
 }
 
 type GetResultsOptions = {

extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts

@@ -1,18 +1,25 @@
 import {
   readQueryResults,
-  runQuery,
+  runExternalApiQueries,
 } from "../../../../src/data-extensions-editor/external-api-usage-queries";
 import { createMockLogger } from "../../../__mocks__/loggerMock";
-import { DatabaseKind } from "../../../../src/databases/local-databases";
+import {
+  DatabaseItem,
+  DatabaseKind,
+} from "../../../../src/databases/local-databases";
 import { dirSync, file } from "tmp-promise";
 import { QueryResultType } from "../../../../src/query-server/new-messages";
 import { fetchExternalApiQueries } from "../../../../src/data-extensions-editor/queries";
 import * as log from "../../../../src/common/logging/notifications";
 import { RedactableError } from "../../../../src/common/errors";
 import { showAndLogExceptionWithTelemetry } from "../../../../src/common/logging";
 import { QueryLanguage } from "../../../../src/common/query-language";
-import { mockedUri } from "../../utils/mocking.helpers";
+import { mockedObject, mockedUri } from "../../utils/mocking.helpers";
 import { Mode } from "../../../../src/data-extensions-editor/shared/mode";
+import { join } from "path";
+import { CodeQLCliServer } from "../../../../src/codeql-cli/cli";
+import { QueryRunner } from "../../../../src/query-server";
+import { QueryOutputDir } from "../../../../src/run-queries-shared";
 
 describe("external api usage query", () => {
   describe("runQuery", () => {
@@ -27,39 +34,44 @@ describe("external api usage query", () => {
         typeof showAndLogExceptionWithTelemetry
       > = jest.spyOn(log, "showAndLogExceptionWithTelemetry");
 
-      const logPath = (await file()).path;
+      const outputDir = new QueryOutputDir(join((await file()).path, "1"));
 
       const query = fetchExternalApiQueries[language];
       if (!query) {
         throw new Error(`No query found for language ${language}`);
       }
 
       const options = {
-        cliServer: {
+        cliServer: mockedObject<CodeQLCliServer>({
           resolveQlpacks: jest.fn().mockResolvedValue({
             "my/extensions": "/a/b/c/",
           }),
-        },
-        queryRunner: {
+          packPacklist: jest
+            .fn()
+            .mockResolvedValue([
+              "/a/b/c/qlpack.yml",
+              "/a/b/c/qlpack.lock.yml",
+              "/a/b/c/qlpack2.yml",
+            ]),
+        }),
+        queryRunner: mockedObject<QueryRunner>({
           createQueryRun: jest.fn().mockReturnValue({
             evaluate: jest.fn().mockResolvedValue({
               resultType: QueryResultType.CANCELLATION,
             }),
-            outputDir: {
-              logPath,
-            },
+            outputDir,
           }),
           logger: createMockLogger(),
-        },
-        databaseItem: {
+        }),
+        databaseItem: mockedObject<DatabaseItem>({
           databaseUri: mockedUri("/a/b/c/src.zip"),
           contents: {
             kind: DatabaseKind.Database,
             name: "foo",
             datasetUri: mockedUri(),
           },
           language,
-        },
+        }),
         queryStorageDir: "/tmp/queries",
         queryDir,
         progress: jest.fn(),
@@ -69,7 +81,9 @@ describe("external api usage query", () => {
         },
       };
 
-      expect(await runQuery(Mode.Application, options)).toBeUndefined();
+      expect(
+        await runExternalApiQueries(Mode.Application, options),
+      ).toBeUndefined();
       expect(showAndLogExceptionWithTelemetrySpy).toHaveBeenCalledWith(
         expect.anything(),
         undefined,
@@ -78,39 +92,48 @@ describe("external api usage query", () => {
     });
 
     it("should run query for random language", async () => {
-      const logPath = (await file()).path;
+      const outputDir = new QueryOutputDir(join((await file()).path, "1"));
 
       const query = fetchExternalApiQueries[language];
       if (!query) {
         throw new Error(`No query found for language ${language}`);
       }
 
       const options = {
-        cliServer: {
+        cliServer: mockedObject<CodeQLCliServer>({
           resolveQlpacks: jest.fn().mockResolvedValue({
             "my/extensions": "/a/b/c/",
           }),
-        },
-        queryRunner: {
+          packPacklist: jest
+            .fn()
+            .mockResolvedValue([
+              "/a/b/c/qlpack.yml",
+              "/a/b/c/qlpack.lock.yml",
+              "/a/b/c/qlpack2.yml",
+            ]),
+          bqrsInfo: jest.fn().mockResolvedValue({
+            "result-sets": [],
+          }),
+        }),
+        queryRunner: mockedObject<QueryRunner>({
           createQueryRun: jest.fn().mockReturnValue({
             evaluate: jest.fn().mockResolvedValue({
               resultType: QueryResultType.SUCCESS,
+              outputDir,
             }),
-            outputDir: {
-              logPath,
-            },
+            outputDir,
           }),
           logger: createMockLogger(),
-        },
-        databaseItem: {
+        }),
+        databaseItem: mockedObject<DatabaseItem>({
           databaseUri: mockedUri("/a/b/c/src.zip"),
           contents: {
             kind: DatabaseKind.Database,
             name: "foo",
             datasetUri: mockedUri(),
           },
           language,
-        },
+        }),
         queryStorageDir: "/tmp/queries",
         queryDir,
         progress: jest.fn(),
@@ -120,9 +143,9 @@ describe("external api usage query", () => {
         },
       };
 
-      const result = await runQuery(Mode.Framework, options);
+      const result = await runExternalApiQueries(Mode.Framework, options);
 
-      expect(result?.resultType).toEqual(QueryResultType.SUCCESS);
+      expect(result).not.toBeUndefined;
 
       expect(options.cliServer.resolveQlpacks).toHaveBeenCalledTimes(1);
       expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith([], true);