Make model editor open progress bar cancellable

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

@@ -7,7 +7,10 @@ import {
 import { CancellationToken } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { DatabaseItem } from "../databases/local-databases";
-import { ProgressCallback } from "../common/vscode/progress";
+import {
+  ProgressCallback,
+  UserCancellationException,
+} from "../common/vscode/progress";
 import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { join } from "path";
@@ -89,6 +92,13 @@ export async function runModelEditorQueries(
   // For a reference of what this should do in the future, see the previous implementation in
   // https://github.com/github/vscode-codeql/blob/089d3566ef0bc67d9b7cc66e8fd6740b31c1c0b0/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts#L33-L72
 
+  if (token.isCancellationRequested) {
+    throw new UserCancellationException(
+      "Run model editor queries cancelled.",
+      true,
+    );
+  }
+
   progress({
     message: "Resolving QL packs",
     step: 1,
@@ -99,6 +109,13 @@ export async function runModelEditorQueries(
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
+  if (token.isCancellationRequested) {
+    throw new UserCancellationException(
+      "Run model editor queries cancelled.",
+      true,
+    );
+  }
+
   progress({
     message: "Resolving query",
     step: 2,
@@ -143,6 +160,13 @@ export async function runModelEditorQueries(
     return;
   }
 
+  if (token.isCancellationRequested) {
+    throw new UserCancellationException(
+      "Run model editor queries cancelled.",
+      true,
+    );
+  }
+
   // Read the results and covert to internal representation
   progress({
     message: "Decoding results",
@@ -159,6 +183,13 @@ export async function runModelEditorQueries(
     return;
   }
 
+  if (token.isCancellationRequested) {
+    throw new UserCancellationException(
+      "Run model editor queries cancelled.",
+      true,
+    );
+  }
+
   progress({
     message: "Finalizing results",
     step: 1950,

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

@@ -1,4 +1,5 @@
 import {
+  CancellationToken,
   CancellationTokenSource,
   Tab,
   TabInputWebview,
@@ -14,7 +15,11 @@ import {
   FromModelEditorMessage,
   ToModelEditorMessage,
 } from "../common/interface-types";
-import { ProgressCallback, withProgress } from "../common/vscode/progress";
+import {
+  ProgressCallback,
+  UserCancellationException,
+  withProgress,
+} from "../common/vscode/progress";
 import { QueryRunner } from "../query-server";
 import {
   showAndLogErrorMessage,
@@ -338,8 +343,8 @@ export class ModelEditorView extends AbstractWebview<
 
     await Promise.all([
       this.setViewState(),
-      withProgress((progress) => this.loadMethods(progress), {
-        cancellable: false,
+      withProgress((progress, token) => this.loadMethods(progress, token), {
+        cancellable: true,
       }),
       this.loadExistingModeledMethods(),
     ]);
@@ -423,11 +428,16 @@ export class ModelEditorView extends AbstractWebview<
     }
   }
 
-  protected async loadMethods(progress: ProgressCallback): Promise<void> {
+  protected async loadMethods(
+    progress: ProgressCallback,
+    token?: CancellationToken,
+  ): Promise<void> {
     const mode = this.modelingStore.getMode(this.databaseItem);
 
     try {
-      const cancellationTokenSource = new CancellationTokenSource();
+      if (!token) {
+        token = new CancellationTokenSource().token;
+      }
       const queryResult = await runModelEditorQueries(mode, {
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
@@ -441,12 +451,19 @@ export class ModelEditorView extends AbstractWebview<
             ...update,
             message: `Loading models: ${update.message}`,
           }),
-        token: cancellationTokenSource.token,
+        token,
       });
       if (!queryResult) {
         return;
       }
 
+      if (token.isCancellationRequested) {
+        throw new UserCancellationException(
+          "Model editor: Load methods cancelled.",
+          true,
+        );
+      }
+
       this.modelingStore.setMethods(this.databaseItem, queryResult);
     } catch (err) {
       void showAndLogExceptionWithTelemetry(
@@ -578,6 +595,7 @@ export class ModelEditorView extends AbstractWebview<
         this.modelConfig,
         this.app.logger,
         progress,
+        token,
         3,
       );
       if (!modelFile) {