Auto-download analyses results (#1098)

Author: charisk

extensions/ql-vscode/src/extension.ts

@@ -80,6 +80,7 @@ import { CodeQlStatusBarHandler } from './status-bar';
 import { Credentials } from './authentication';
 import { RemoteQueriesManager } from './remote-queries/remote-queries-manager';
 import { RemoteQuery } from './remote-queries/remote-query';
+import { RemoteQueryResult } from './remote-queries/remote-query-result';
 import { URLSearchParams } from 'url';
 import { RemoteQueriesInterfaceManager } from './remote-queries/remote-queries-interface';
 import { sampleRemoteQuery, sampleRemoteQueryResult } from './remote-queries/sample-data';
@@ -778,7 +779,7 @@ async function activateWithInstalledDistribution(
   );
 
   void logger.log('Initializing remote queries interface.');
-  const rqm = new RemoteQueriesManager(ctx, logger, cliServer);
+  const rqm = new RemoteQueriesManager(ctx, cliServer, logger);
 
   registerRemoteQueryTextProvider();
 
@@ -816,6 +817,13 @@ async function activateWithInstalledDistribution(
       await rqm.monitorRemoteQuery(query, token);
     }));
 
+  ctx.subscriptions.push(
+    commandRunner('codeQL.autoDownloadRemoteQueryResults', async (
+      queryResult: RemoteQueryResult,
+      token: CancellationToken) => {
+      await rqm.autoDownloadRemoteQueryResults(queryResult, token);
+    }));
+
   ctx.subscriptions.push(
     commandRunner('codeQL.showFakeRemoteQueryResults', async () => {
       const analysisResultsManager = new AnalysesResultsManager(ctx, logger);

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

@@ -1,4 +1,4 @@
-import { ExtensionContext } from 'vscode';
+import { CancellationToken, ExtensionContext } from 'vscode';
 import { Credentials } from '../authentication';
 import { Logger } from '../logging';
 import { downloadArtifactFromLink } from './gh-actions-api-client';
@@ -7,6 +7,7 @@ import * as fs from 'fs-extra';
 import { AnalysisSummary } from './shared/remote-query-result';
 import * as sarif from 'sarif';
 import { AnalysisResults, QueryResult } from './shared/analysis-result';
+import { UserCancellationException } from '../commandRunner';
 
 export class AnalysesResultsManager {
   // Store for the results of various analyses for a single remote query.
@@ -21,6 +22,7 @@ export class AnalysesResultsManager {
 
   public async downloadAnalysisResults(
     analysisSummary: AnalysisSummary,
+    publishResults: (analysesResults: AnalysisResults[]) => Promise<void>
   ): Promise<void> {
     if (this.analysesResults.some(x => x.nwo === analysisSummary.nwo)) {
       // We already have the results for this analysis, don't download again.
@@ -32,17 +34,35 @@ export class AnalysesResultsManager {
     void this.logger.log(`Downloading and processing results for ${analysisSummary.nwo}`);
 
     await this.downloadSingleAnalysisResults(analysisSummary, credentials);
+    await publishResults(this.analysesResults);
   }
 
-  public async downloadAllResults(
-    analysisSummaries: AnalysisSummary[],
+  public async downloadAnalysesResults(
+    analysesToDownload: AnalysisSummary[],
+    token: CancellationToken | undefined,
+    publishResults: (analysesResults: AnalysisResults[]) => Promise<void>
   ): Promise<void> {
     const credentials = await Credentials.initialize(this.ctx);
 
-    void this.logger.log('Downloading and processing all results');
+    void this.logger.log('Downloading and processing analyses results');
 
-    for (const analysis of analysisSummaries) {
-      await this.downloadSingleAnalysisResults(analysis, credentials);
+    const batchSize = 3;
+    const numOfBatches = Math.ceil(analysesToDownload.length / batchSize);
+
+    for (let i = 0; i < analysesToDownload.length; i += batchSize) {
+      if (token?.isCancellationRequested) {
+        throw new UserCancellationException('Downloading of analyses results has been cancelled', true);
+      }
+
+      const batch = analysesToDownload.slice(i, i + batchSize);
+      const batchTasks = batch.map(analysis => this.downloadSingleAnalysisResults(analysis, credentials));
+
+      const nwos = batch.map(a => a.nwo).join(', ');
+      void this.logger.log(`Downloading batch ${Math.floor(i / batchSize) + 1} of ${numOfBatches} (${nwos})`);
+
+      await Promise.all(batchTasks);
+
+      await publishResults(this.analysesResults);
     }
   }
 

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

@@ -51,6 +51,8 @@ export class RemoteQueriesInterfaceManager {
       t: 'setRemoteQueryResult',
       queryResult: this.buildViewModel(query, queryResult)
     });
+
+    await this.setAnalysisResults(this.analysesResultsManager.getAnalysesResults());
   }
 
   /**
@@ -201,20 +203,25 @@ export class RemoteQueriesInterfaceManager {
   }
 
   private async downloadAnalysisResults(msg: RemoteQueryDownloadAnalysisResultsMessage): Promise<void> {
-    await this.analysesResultsManager.downloadAnalysisResults(msg.analysisSummary);
-    await this.setAnalysisResults(this.analysesResultsManager.getAnalysesResults());
+    await this.analysesResultsManager.downloadAnalysisResults(
+      msg.analysisSummary,
+      results => this.setAnalysisResults(results));
   }
 
   private async downloadAllAnalysesResults(msg: RemoteQueryDownloadAllAnalysesResultsMessage): Promise<void> {
-    await this.analysesResultsManager.downloadAllResults(msg.analysisSummaries);
-    await this.setAnalysisResults(this.analysesResultsManager.getAnalysesResults());
+    await this.analysesResultsManager.downloadAnalysesResults(
+      msg.analysisSummaries,
+      undefined,
+      results => this.setAnalysisResults(results));
   }
 
-  private async setAnalysisResults(analysesResults: AnalysisResults[]): Promise<void> {
-    await this.postMessage({
-      t: 'setAnalysesResults',
-      analysesResults: analysesResults
-    });
+  public async setAnalysisResults(analysesResults: AnalysisResults[]): Promise<void> {
+    if (this.panel?.active) {
+      await this.postMessage({
+        t: 'setAnalysesResults',
+        analysesResults: analysesResults
+      });
+    }
   }
 
   private postMessage(msg: ToRemoteQueriesMessage): Thenable<boolean> {

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

@@ -14,16 +14,21 @@ import { RemoteQueryResult } from './remote-query-result';
 import { DownloadLink } from './download-link';
 import { AnalysesResultsManager } from './analyses-results-manager';
 
+const autoDownloadMaxSize = 300 * 1024;
+const autoDownloadMaxCount = 100;
+
 export class RemoteQueriesManager {
   private readonly remoteQueriesMonitor: RemoteQueriesMonitor;
   private readonly analysesResultsManager: AnalysesResultsManager;
+  private readonly interfaceManager: RemoteQueriesInterfaceManager;
 
   constructor(
     private readonly ctx: ExtensionContext,
-    private readonly logger: Logger,
-    private readonly cliServer: CodeQLCliServer
+    private readonly cliServer: CodeQLCliServer,
+    logger: Logger,
   ) {
     this.analysesResultsManager = new AnalysesResultsManager(ctx, logger);
+    this.interfaceManager = new RemoteQueriesInterfaceManager(ctx, logger, this.analysesResultsManager);
     this.remoteQueriesMonitor = new RemoteQueriesMonitor(ctx, logger);
   }
 
@@ -65,13 +70,16 @@ export class RemoteQueriesManager {
 
       const queryResult = this.mapQueryResult(executionEndTime, resultIndex);
 
+      // Kick off auto-download of results.
+      void commands.executeCommand('codeQL.autoDownloadRemoteQueryResults', queryResult);
+
       const totalResultCount = queryResult.analysisSummaries.reduce((acc, cur) => acc + cur.resultCount, 0);
       const message = `Query "${query.queryName}" run on ${query.repositories.length} repositories and returned ${totalResultCount} results`;
 
       const shouldOpenView = await showInformationMessageWithAction(message, 'View');
       if (shouldOpenView) {
-        const rqim = new RemoteQueriesInterfaceManager(this.ctx, this.logger, this.analysesResultsManager);
-        await rqim.showResults(query, queryResult);
+        await this.interfaceManager.showResults(query, queryResult);
+
       }
     } else if (queryResult.status === 'CompletedUnsuccessfully') {
       await showAndLogErrorMessage(`Remote query execution failed. Error: ${queryResult.error}`);
@@ -81,6 +89,26 @@ export class RemoteQueriesManager {
     }
   }
 
+  public async autoDownloadRemoteQueryResults(
+    queryResult: RemoteQueryResult,
+    token: CancellationToken
+  ): Promise<void> {
+    const analysesToDownload = queryResult.analysisSummaries
+      .filter(a => a.fileSizeInBytes < autoDownloadMaxSize)
+      .slice(0, autoDownloadMaxCount)
+      .map(a => ({
+        nwo: a.nwo,
+        resultCount: a.resultCount,
+        downloadLink: a.downloadLink,
+        fileSize: String(a.fileSizeInBytes)
+      }));
+
+    await this.analysesResultsManager.downloadAnalysesResults(
+      analysesToDownload,
+      token,
+      results => this.interfaceManager.setAnalysisResults(results));
+  }
+
   private mapQueryResult(executionEndTime: Date, resultIndex: RemoteQueryResultIndex): RemoteQueryResult {
     const analysisSummaries = resultIndex.items.map(item => ({
       nwo: item.nwo,