Merge pull request #3454 from github/koesie10/cancellable-variant-analysis

Make the "Run variant analysis against published pack" command cancellable

Author: koesie10

extensions/ql-vscode/src/codeql-cli/cli.ts

@@ -1,5 +1,5 @@
 import type { CodeQLCliServer } from "./cli";
-import type { Uri } from "vscode";
+import type { CancellationToken, Uri } from "vscode";
 import { window } from "vscode";
 import {
   getLanguageDisplayName,
@@ -50,6 +50,7 @@ export async function findLanguage(
 export async function askForLanguage(
   cliServer: CodeQLCliServer,
   throwOnEmpty = true,
+  token?: CancellationToken,
 ): Promise<QueryLanguage | undefined> {
   const supportedLanguages = await cliServer.getSupportedLanguages();
 
@@ -62,10 +63,14 @@ export async function askForLanguage(
     }))
     .sort((a, b) => a.label.localeCompare(b.label));
 
-  const selectedItem = await window.showQuickPick(items, {
-    placeHolder: "Select target query language",
-    ignoreFocusOut: true,
-  });
+  const selectedItem = await window.showQuickPick(
+    items,
+    {
+      placeHolder: "Select target query language",
+      ignoreFocusOut: true,
+    },
+    token,
+  );
   if (!selectedItem) {
     // This only happens if the user cancels the quick pick.
     if (throwOnEmpty) {

extensions/ql-vscode/src/codeql-cli/query-language.ts

@@ -58,6 +58,7 @@ export class ModelEvaluator extends DisposableObject {
       this.app.logger,
       this.cliServer,
       this.language,
+      this.cancellationSource.token,
     );
 
     if (!qlPack) {

extensions/ql-vscode/src/model-editor/model-evaluator.ts

@@ -9,16 +9,18 @@ import type { QlPackDetails } from "./ql-pack-details";
 import { getQlPackFilePath } from "../common/ql";
 import type { SuiteInstruction } from "../packaging/suite-instruction";
 import { SARIF_RESULTS_QUERY_KINDS } from "../common/query-metadata";
+import type { CancellationToken } from "vscode";
 
 export async function resolveCodeScanningQueryPack(
   logger: BaseLogger,
   cliServer: CodeQLCliServer,
   language: QueryLanguage,
+  token: CancellationToken,
 ): Promise<QlPackDetails> {
   // Get pack
   void logger.log(`Downloading pack for language: ${language}`);
   const packName = `codeql/${language}-queries`;
-  const packDownloadResult = await cliServer.packDownload([packName]);
+  const packDownloadResult = await cliServer.packDownload([packName], token);
   const downloadedPack = packDownloadResult.packs[0];
 
   const packDir = join(

extensions/ql-vscode/src/variant-analysis/code-scanning-pack.ts

@@ -54,6 +54,7 @@ async function generateQueryPack(
   cliServer: CodeQLCliServer,
   qlPackDetails: QlPackDetails,
   tmpDir: RemoteQueryTempDir,
+  token: CancellationToken,
 ): Promise<string> {
   const workspaceFolders = getOnDiskWorkspaceFolders();
   const extensionPacks = await getExtensionPacksToInject(
@@ -148,6 +149,7 @@ async function generateQueryPack(
     bundlePath,
     tmpDir.compiledPackDir,
     precompilationOpts,
+    token,
   );
   const base64Pack = (await readFile(bundlePath)).toString("base64");
   return base64Pack;
@@ -331,7 +333,12 @@ export async function prepareRemoteQueryRun(
   let base64Pack: string;
 
   try {
-    base64Pack = await generateQueryPack(cliServer, qlPackDetails, tempDir);
+    base64Pack = await generateQueryPack(
+      cliServer,
+      qlPackDetails,
+      tempDir,
+      token,
+    );
   } finally {
     await tempDir.remoteQueryDir.cleanup();
   }

extensions/ql-vscode/src/variant-analysis/run-remote-query.ts

@@ -221,42 +221,49 @@ export class VariantAnalysisManager
   }
 
   public async runVariantAnalysisFromPublishedPack(): Promise<void> {
-    return withProgress(async (progress, token) => {
-      progress({
-        maxStep: 7,
-        step: 0,
-        message: "Determining query language",
-      });
-
-      const language = await askForLanguage(this.cliServer);
-      if (!language) {
-        return;
-      }
-
-      progress({
-        maxStep: 7,
-        step: 2,
-        message: "Downloading query pack and resolving queries",
-      });
+    return withProgress(
+      async (progress, token) => {
+        progress({
+          maxStep: 7,
+          step: 0,
+          message: "Determining query language",
+        });
+
+        const language = await askForLanguage(this.cliServer, true, token);
+        if (!language) {
+          return;
+        }
 
-      // Build up details to pass to the functions that run the variant analysis.
-      const qlPackDetails = await resolveCodeScanningQueryPack(
-        this.app.logger,
-        this.cliServer,
-        language,
-      );
+        progress({
+          maxStep: 7,
+          step: 2,
+          message: "Downloading query pack and resolving queries",
+        });
+
+        // Build up details to pass to the functions that run the variant analysis.
+        const qlPackDetails = await resolveCodeScanningQueryPack(
+          this.app.logger,
+          this.cliServer,
+          language,
+          token,
+        );
 
-      await this.runVariantAnalysis(
-        qlPackDetails,
-        (p) =>
-          progress({
-            ...p,
-            maxStep: p.maxStep + 3,
-            step: p.step + 3,
-          }),
-        token,
-      );
-    });
+        await this.runVariantAnalysis(
+          qlPackDetails,
+          (p) =>
+            progress({
+              ...p,
+              maxStep: p.maxStep + 3,
+              step: p.step + 3,
+            }),
+          token,
+        );
+      },
+      {
+        title: "Run Variant Analysis",
+        cancellable: true,
+      },
+    );
   }
 
   private async runVariantAnalysisCommand(queryFiles: Uri[]): Promise<void> {

extensions/ql-vscode/src/variant-analysis/variant-analysis-manager.ts

@@ -9,6 +9,8 @@ const config: Config = {
   // CLI integration tests call into the CLI and execute queries, so these are expected to take a lot longer
   // than the default 5 seconds.
   testTimeout: 180_000, // 3 minutes
+  // Ensure that Jest exits even when there are some remaining handles open
+  forceExit: true,
 };
 
 export default config;

extensions/ql-vscode/test/vscode-tests/cli-integration/jest.config.ts

@@ -1,3 +1,4 @@
+import { CancellationTokenSource } from "vscode";
 import type { CodeQLCliServer } from "../../../../src/codeql-cli/cli";
 import type { App } from "../../../../src/common/app";
 import { QueryLanguage } from "../../../../src/common/query-language";
@@ -20,6 +21,7 @@ describe("Code Scanning pack", () => {
       app.logger,
       cli,
       QueryLanguage.Javascript,
+      new CancellationTokenSource().token,
     );
     // Should include queries. Just check that at least one known query exists.
     // It doesn't particularly matter which query we check for.