Merge pull request #1291 from github/aeisenberg/handle-remote-cancel

Handle cancelling of remote queries

Author: aeisenberg

extensions/ql-vscode/src/authentication.ts

@@ -7,7 +7,7 @@ const GITHUB_AUTH_PROVIDER_ID = 'github';
 // https://docs.github.com/apps/building-oauth-apps/understanding-scopes-for-oauth-apps
 const SCOPES = ['repo'];
 
-/** 
+/**
  * Handles authentication to GitHub, using the VS Code [authentication API](https://code.visualstudio.com/api/references/vscode-api#authentication).
  */
 export class Credentials {
@@ -18,6 +18,15 @@ export class Credentials {
   // eslint-disable-next-line @typescript-eslint/no-empty-function
   private constructor() { }
 
+  /**
+   * Initializes an instance of credentials with an octokit instance.
+   *
+   * Do not call this method until you know you actually need an instance of credentials.
+   * since calling this method will require the user to log in.
+   *
+   * @param context The extension context.
+   * @returns An instance of credentials.
+   */
   static async initialize(context: vscode.ExtensionContext): Promise<Credentials> {
     const c = new Credentials();
     c.registerListeners(context);

extensions/ql-vscode/src/query-history.ts

@@ -36,6 +36,8 @@ import { QueryStatus } from './query-status';
 import { slurpQueryHistory, splatQueryHistory } from './query-serialization';
 import * as fs from 'fs-extra';
 import { CliVersionConstraint } from './cli';
+import { Credentials } from './authentication';
+import { cancelRemoteQuery } from './remote-queries/gh-actions-api-client';
 
 /**
  * query-history.ts
@@ -316,7 +318,7 @@ export class QueryHistoryManager extends DisposableObject {
     private qs: QueryServerClient,
     private dbm: DatabaseManager,
     private queryStorageDir: string,
-    ctx: ExtensionContext,
+    private ctx: ExtensionContext,
     private queryHistoryConfigListener: QueryHistoryConfig,
     private doCompareCallback: (
       from: CompletedLocalQueryInfo,
@@ -512,6 +514,10 @@ export class QueryHistoryManager extends DisposableObject {
     this.registerQueryHistoryScrubber(queryHistoryConfigListener, ctx);
   }
 
+  private getCredentials() {
+    return Credentials.initialize(this.ctx);
+  }
+
   /**
    * Register and create the history scrubber.
    */
@@ -816,7 +822,7 @@ export class QueryHistoryManager extends DisposableObject {
     }
 
     if (finalSingleItem.evalLogSummaryLocation) {
-        await this.tryOpenExternalFile(finalSingleItem.evalLogSummaryLocation);
+      await this.tryOpenExternalFile(finalSingleItem.evalLogSummaryLocation);
     } else {
       this.warnNoEvalLogSummary();
     }
@@ -830,11 +836,20 @@ export class QueryHistoryManager extends DisposableObject {
     // In the future, we may support cancelling remote queries, but this is not a short term plan.
     const { finalSingleItem, finalMultiSelect } = this.determineSelection(singleItem, multiSelect);
 
-    (finalMultiSelect || [finalSingleItem]).forEach((item) => {
-      if (item.status === QueryStatus.InProgress && item.t === 'local') {
-        item.cancel();
+    const selected = finalMultiSelect || [finalSingleItem];
+    const results = selected.map(async item => {
+      if (item.status === QueryStatus.InProgress) {
+        if (item.t === 'local') {
+          item.cancel();
+        } else if (item.t === 'remote') {
+          void showAndLogInformationMessage('Cancelling variant analysis. This may take a while.');
+          const credentials = await this.getCredentials();
+          await cancelRemoteQuery(credentials, item.remoteQuery);
+        }
       }
     });
+
+    await Promise.all(results);
   }
 
   async handleShowQueryText(

extensions/ql-vscode/src/remote-queries/gh-actions-api-client.ts

@@ -74,6 +74,18 @@ export async function getRemoteQueryIndex(
   };
 }
 
+export async function cancelRemoteQuery(
+  credentials: Credentials,
+  remoteQuery: RemoteQuery
+): Promise<void> {
+  const octokit = await credentials.getOctokit();
+  const { actionsWorkflowRunId, controllerRepository: { owner, name } } = remoteQuery;
+  const response = await octokit.request(`POST /repos/${owner}/${name}/actions/runs/${actionsWorkflowRunId}/cancel`);
+  if (response.status >= 300) {
+    throw new Error(`Error cancelling variant analysis: ${response.status} ${response?.data?.message || ''}`);
+  }
+}
+
 export async function downloadArtifactFromLink(
   credentials: Credentials,
   storagePath: string,

extensions/ql-vscode/src/remote-queries/remote-queries-manager.ts

@@ -6,7 +6,7 @@ import * as fs from 'fs-extra';
 import { Credentials } from '../authentication';
 import { CodeQLCliServer } from '../cli';
 import { ProgressCallback } from '../commandRunner';
-import { createTimestampFile, showAndLogErrorMessage, showInformationMessageWithAction } from '../helpers';
+import { createTimestampFile, showAndLogErrorMessage, showAndLogInformationMessage, showInformationMessageWithAction } from '../helpers';
 import { Logger } from '../logging';
 import { runRemoteQuery } from './run-remote-query';
 import { RemoteQueriesInterfaceManager } from './remote-queries-interface';
@@ -155,13 +155,20 @@ export class RemoteQueriesManager extends DisposableObject {
         queryItem.status = QueryStatus.Failed;
       }
     } else if (queryWorkflowResult.status === 'CompletedUnsuccessfully') {
-      queryItem.failureReason = queryWorkflowResult.error;
-      queryItem.status = QueryStatus.Failed;
-      void showAndLogErrorMessage(`Variant analysis execution failed. Error: ${queryWorkflowResult.error}`);
+      if (queryWorkflowResult.error?.includes('cancelled')) {
+        // workflow was cancelled on the server
+        queryItem.failureReason = 'Cancelled';
+        queryItem.status = QueryStatus.Failed;
+        void showAndLogInformationMessage('Variant analysis was cancelled');
+      } else {
+        queryItem.failureReason = queryWorkflowResult.error;
+        queryItem.status = QueryStatus.Failed;
+        void showAndLogErrorMessage(`Variant analysis execution failed. Error: ${queryWorkflowResult.error}`);
+      }
     } else if (queryWorkflowResult.status === 'Cancelled') {
       queryItem.failureReason = 'Cancelled';
       queryItem.status = QueryStatus.Failed;
-      void showAndLogErrorMessage('Variant analysis monitoring was cancelled');
+      void showAndLogErrorMessage('Variant analysis was cancelled');
     } else if (queryWorkflowResult.status === 'InProgress') {
       // Should not get here. Only including this to ensure `assertNever` uses proper type checking.
       void showAndLogErrorMessage(`Unexpected status: ${queryWorkflowResult.status}`);

extensions/ql-vscode/src/vscode-tests/no-workspace/remote-queries/gh-actions-api-client.test.ts

@@ -0,0 +1,52 @@
+import { expect } from 'chai';
+import * as sinon from 'sinon';
+import { Credentials } from '../../../authentication';
+import { cancelRemoteQuery } from '../../../remote-queries/gh-actions-api-client';
+import { RemoteQuery } from '../../../remote-queries/remote-query';
+
+describe('gh-actions-api-client', () => {
+  let sandbox: sinon.SinonSandbox;
+  let mockCredentials: Credentials;
+  let mockResponse: sinon.SinonStub<any, Promise<{ status: number }>>;
+
+  beforeEach(() => {
+    sandbox = sinon.createSandbox();
+    mockCredentials = {
+      getOctokit: () => Promise.resolve({
+        request: mockResponse
+      })
+    } as unknown as Credentials;
+  });
+
+  afterEach(() => {
+    sandbox.restore();
+  });
+
+  describe('cancelRemoteQuery', () => {
+    it('should cancel a remote query', async () => {
+      mockResponse = sinon.stub().resolves({ status: 202 });
+      await cancelRemoteQuery(mockCredentials, createMockRemoteQuery());
+
+      expect(mockResponse.calledOnce).to.be.true;
+      expect(mockResponse.firstCall.args[0]).to.equal('POST /repos/github/codeql/actions/runs/123/cancel');
+    });
+
+    it('should fail to cancel a remote query', async () => {
+      mockResponse = sinon.stub().resolves({ status: 409, data: { message: 'Uh oh!' } });
+
+      await expect(cancelRemoteQuery(mockCredentials, createMockRemoteQuery())).to.be.rejectedWith(/Error cancelling variant analysis: 409 Uh oh!/);
+      expect(mockResponse.calledOnce).to.be.true;
+      expect(mockResponse.firstCall.args[0]).to.equal('POST /repos/github/codeql/actions/runs/123/cancel');
+    });
+
+    function createMockRemoteQuery(): RemoteQuery {
+      return {
+        actionsWorkflowRunId: 123,
+        controllerRepository: {
+          owner: 'github',
+          name: 'codeql'
+        }
+      } as unknown as RemoteQuery;
+    }
+  });
+});