Starcke/automodel queries refactor

extensions/ql-vscode/src/data-extensions-editor/auto-model-codeml-queries.ts

@@ -1,42 +1,21 @@
 import { CodeQLCliServer, SourceInfo } from "../codeql-cli/cli";
-import { QueryRunner } from "../query-server";
+import { CoreCompletedQuery, QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { ProgressCallback } from "../common/vscode/progress";
 import * as Sarif from "sarif";
 import { qlpackOfDatabase, resolveQueries } from "../local-queries";
-import { extLogger } from "../common/logging/vscode";
 import { Mode } from "./shared/mode";
-import { QlPacksForLanguage } from "../databases/qlpack";
-import { createLockFileForStandardQuery } from "../local-queries/standard-queries";
-import { CancellationToken, CancellationTokenSource } from "vscode";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
-import { showAndLogExceptionWithTelemetry, TeeLogger } from "../common/logging";
-import { QueryResultType } from "../query-server/new-messages";
-import { telemetryListener } from "../common/vscode/telemetry";
-import { redactableError } from "../common/errors";
 import { interpretResultsSarif } from "../query-results";
 import { join } from "path";
 import { assertNever } from "../common/helpers-pure";
 import { dir } from "tmp-promise";
 import { writeFile, outputFile } from "fs-extra";
 import { dump as dumpYaml } from "js-yaml";
 import { MethodSignature } from "./external-api-usage";
-
-type AutoModelQueryOptions = {
-  queryTag: string;
-  mode: Mode;
-  cliServer: CodeQLCliServer;
-  queryRunner: QueryRunner;
-  databaseItem: DatabaseItem;
-  qlpack: QlPacksForLanguage;
-  sourceInfo: SourceInfo | undefined;
-  additionalPacks: string[];
-  extensionPacks: string[];
-  queryStorageDir: string;
-
-  progress: ProgressCallback;
-  token: CancellationToken;
-};
+import { runQuery } from "../local-queries/run-query";
+import { QueryMetadata } from "../common/interface-types";
+import { CancellationTokenSource } from "vscode";
 
 function modeTag(mode: Mode): string {
   switch (mode) {
@@ -49,108 +28,6 @@ function modeTag(mode: Mode): string {
   }
 }
 
-async function runAutoModelQuery({
-  queryTag,
-  mode,
-  cliServer,
-  queryRunner,
-  databaseItem,
-  qlpack,
-  sourceInfo,
-  additionalPacks,
-  extensionPacks,
-  queryStorageDir,
-  progress,
-  token,
-}: AutoModelQueryOptions): Promise<Sarif.Log | undefined> {
-  // First, resolve the query that we want to run.
-  // All queries are tagged like this:
-  // internal extract automodel <mode> <queryTag>
-  // Example: internal extract automodel framework-mode candidates
-  const queries = await resolveQueries(
-    cliServer,
-    qlpack,
-    `Extract automodel ${queryTag}`,
-    {
-      kind: "problem",
-      "tags contain all": ["automodel", modeTag(mode), ...queryTag.split(" ")],
-    },
-  );
-  if (queries.length > 1) {
-    throw new Error(
-      `Found multiple auto model queries for ${mode} ${queryTag}. Can't continue`,
-    );
-  }
-  if (queries.length === 0) {
-    throw new Error(
-      `Did not found any auto model queries for ${mode} ${queryTag}. Can't continue`,
-    );
-  }
-
-  const queryPath = queries[0];
-  const { cleanup: cleanupLockFile } = await createLockFileForStandardQuery(
-    cliServer,
-    queryPath,
-  );
-
-  // Get metadata for the query. This is required to interpret the results. We already know the kind is problem
-  // (because of the constraint in resolveQueries), so we don't need any more checks on the metadata.
-  const metadata = await cliServer.resolveMetadata(queryPath);
-
-  const queryRun = queryRunner.createQueryRun(
-    databaseItem.databaseUri.fsPath,
-    {
-      queryPath,
-      quickEvalPosition: undefined,
-      quickEvalCountOnly: false,
-    },
-    false,
-    additionalPacks,
-    extensionPacks,
-    queryStorageDir,
-    undefined,
-    undefined,
-  );
-
-  const completedQuery = await queryRun.evaluate(
-    progress,
-    token,
-    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
-  );
-
-  await cleanupLockFile?.();
-
-  if (completedQuery.resultType !== QueryResultType.SUCCESS) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`Auto-model query ${queryTag} failed: ${
-        completedQuery.message ?? "No message"
-      }`,
-    );
-    return;
-  }
-
-  const interpretedResultsPath = join(
-    queryStorageDir,
-    `interpreted-results-${queryTag.replaceAll(" ", "-")}-${queryRun.id}.sarif`,
-  );
-
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars -- We only need the actual SARIF data, not the extra fields added by SarifInterpretationData
-  const { t, sortState, ...sarif } = await interpretResultsSarif(
-    cliServer,
-    metadata,
-    {
-      resultsPath: completedQuery.outputDir.bqrsPath,
-      interpretedResultsPath,
-    },
-    sourceInfo,
-    ["--sarif-add-snippets"],
-  );
-
-  return sarif;
-}
-
 type AutoModelQueriesOptions = {
   mode: Mode;
   candidateMethods: MethodSignature[];
@@ -177,22 +54,13 @@ export async function runAutoModelQueries({
   progress,
   cancellationTokenSource,
 }: AutoModelQueriesOptions): Promise<AutoModelQueriesResult | undefined> {
-  const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
-
-  // CodeQL needs to have access to the database to be able to retrieve the
-  // snippets from it. The source location prefix is used to determine the
-  // base path of the database.
-  const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
+  // First, resolve the query that we want to run.
+  const queryPath = await resolveAutomodelQuery(
     cliServer,
+    databaseItem,
+    "candidates",
+    mode,
   );
-  const sourceArchiveUri = databaseItem.sourceArchive;
-  const sourceInfo =
-    sourceArchiveUri === undefined
-      ? undefined
-      : {
-          sourceArchive: sourceArchiveUri.fsPath,
-          sourceLocationPrefix,
-        };
 
   // Generate a pack containing the candidate filters
   const filterPackDir = await generateCandidateFilterPack(
@@ -205,30 +73,89 @@ export async function runAutoModelQueries({
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
-  const candidates = await runAutoModelQuery({
-    mode,
-    queryTag: "candidates",
+  // Run the actual query
+  const completedQuery = await runQuery({
     cliServer,
     queryRunner,
     databaseItem,
-    qlpack,
-    sourceInfo,
+    queryPath,
+    queryStorageDir,
     additionalPacks,
     extensionPacks,
-    queryStorageDir,
     progress,
     token: cancellationTokenSource.token,
   });
 
-  if (!candidates) {
+  if (!completedQuery) {
     return undefined;
   }
 
+  // Get metadata for the query. This is required to interpret the results. We already know the kind is problem
+  // (because of the constraint in resolveQueries), so we don't need any more checks on the metadata.
+  const metadata = await cliServer.resolveMetadata(queryPath);
+
+  // CodeQL needs to have access to the database to be able to retrieve the
+  // snippets from it. The source location prefix is used to determine the
+  // base path of the database.
+  const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
+    cliServer,
+  );
+  const sourceArchiveUri = databaseItem.sourceArchive;
+  const sourceInfo =
+    sourceArchiveUri === undefined
+      ? undefined
+      : {
+          sourceArchive: sourceArchiveUri.fsPath,
+          sourceLocationPrefix,
+        };
+
+  const candidates = await interpretAutomodelResults(
+    cliServer,
+    completedQuery,
+    metadata,
+    sourceInfo,
+  );
+
   return {
     candidates,
   };
 }
 
+async function resolveAutomodelQuery(
+  cliServer: CodeQLCliServer,
+  databaseItem: DatabaseItem,
+  queryTag: string,
+  mode: Mode,
+): Promise<string> {
+  const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
+
+  // First, resolve the query that we want to run.
+  // All queries are tagged like this:
+  // internal extract automodel <mode> <queryTag>
+  // Example: internal extract automodel framework-mode candidates
+  const queries = await resolveQueries(
+    cliServer,
+    qlpack,
+    `Extract automodel ${queryTag}`,
+    {
+      kind: "problem",
+      "tags contain all": ["automodel", modeTag(mode), ...queryTag.split(" ")],
+    },
+  );
+  if (queries.length > 1) {
+    throw new Error(
+      `Found multiple auto model queries for ${mode} ${queryTag}. Can't continue`,
+    );
+  }
+  if (queries.length === 0) {
+    throw new Error(
+      `Did not found any auto model queries for ${mode} ${queryTag}. Can't continue`,
+    );
+  }
+
+  return queries[0];
+}
+
 /**
  * generateCandidateFilterPack will create a temporary extension pack.
  * This pack will contain a filter that will restrict the automodel queries
@@ -284,3 +211,28 @@ export async function generateCandidateFilterPack(
 
   return packDir;
 }
+
+async function interpretAutomodelResults(
+  cliServer: CodeQLCliServer,
+  completedQuery: CoreCompletedQuery,
+  metadata: QueryMetadata,
+  sourceInfo: SourceInfo | undefined,
+): Promise<Sarif.Log> {
+  const interpretedResultsPath = join(
+    completedQuery.outputDir.querySaveDir,
+    "results.sarif",
+  );
+
+  const { ...sarif } = await interpretResultsSarif(
+    cliServer,
+    metadata,
+    {
+      resultsPath: completedQuery.outputDir.bqrsPath,
+      interpretedResultsPath,
+    },
+    sourceInfo,
+    ["--sarif-add-snippets"],
+  );
+
+  return sarif;
+}

extensions/ql-vscode/src/local-queries/run-query.ts

@@ -0,0 +1,78 @@
+import { CancellationToken } from "vscode";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { ProgressCallback } from "../common/vscode/progress";
+import { DatabaseItem } from "../databases/local-databases";
+import { CoreCompletedQuery, QueryRunner } from "../query-server";
+import { createLockFileForStandardQuery } from "./standard-queries";
+import { TeeLogger, showAndLogExceptionWithTelemetry } from "../common/logging";
+import { QueryResultType } from "../query-server/new-messages";
+import { extLogger } from "../common/logging/vscode";
+import { telemetryListener } from "../common/vscode/telemetry";
+import { redactableError } from "../common/errors";
+import { basename } from "path";
+
+type RunQueryOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  databaseItem: DatabaseItem;
+  queryPath: string;
+  queryStorageDir: string;
+  additionalPacks: string[];
+  extensionPacks: string[] | undefined;
+  progress: ProgressCallback;
+  token: CancellationToken;
+};
+
+export async function runQuery({
+  cliServer,
+  queryRunner,
+  databaseItem,
+  queryPath,
+  queryStorageDir,
+  additionalPacks,
+  extensionPacks,
+  progress,
+  token,
+}: RunQueryOptions): Promise<CoreCompletedQuery | undefined> {
+  // Create a lock file for the query. This is required to resolve dependencies and library path for the query.
+  const { cleanup: cleanupLockFile } = await createLockFileForStandardQuery(
+    cliServer,
+    queryPath,
+  );
+
+  // Create a query run to execute
+  const queryRun = queryRunner.createQueryRun(
+    databaseItem.databaseUri.fsPath,
+    {
+      queryPath,
+      quickEvalPosition: undefined,
+      quickEvalCountOnly: false,
+    },
+    false,
+    additionalPacks,
+    extensionPacks,
+    queryStorageDir,
+    undefined,
+    undefined,
+  );
+
+  const completedQuery = await queryRun.evaluate(
+    progress,
+    token,
+    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
+  );
+
+  await cleanupLockFile?.();
+
+  if (completedQuery.resultType !== QueryResultType.SUCCESS) {
+    void showAndLogExceptionWithTelemetry(
+      extLogger,
+      telemetryListener,
+      redactableError`Failed to run ${basename(queryPath)} query: ${
+        completedQuery.message ?? "No message"
+      }`,
+    );
+    return;
+  }
+  return completedQuery;
+}