Allow exporting of results for non-alert queries

github · Aug 19, 2021 · 8c1e8a4 · 8c1e8a4
1 parent 089b23f
commit 8c1e8a4
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 3 deletions.
diff --git a/extensions/ql-vscode/CHANGELOG.md b/extensions/ql-vscode/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## [UNRELEASED]
 
 - Add support for filename pattern in history view. [#930](https://github.com/github/vscode-codeql/pull/930)
+- Add an option _Export Results (CSV)_ to export the results of a non-alert query. The existing options for alert queries have been renamed to _View Alerts_ to avoid confusion. [#929](https://github.com/github/vscode-codeql/pull/929)
 
 ## 1.5.3 - 18 August 2021
 

diff --git a/extensions/ql-vscode/package.json b/extensions/ql-vscode/package.json
@@ -461,13 +461,17 @@
         "command": "codeQLQueryHistory.showQueryText",
         "title": "Show Query Text"
       },
+      {
+        "command": "codeQLQueryHistory.exportCsvResults",
+        "title": "Export Results (CSV)"
+      },
       {
         "command": "codeQLQueryHistory.viewCsvResults",
-        "title": "View Results (CSV)"
+        "title": "View Alerts (CSV)"
       },
       {
         "command": "codeQLQueryHistory.viewSarifResults",
-        "title": "View Results (SARIF)"
+        "title": "View Alerts (SARIF)"
       },
       {
         "command": "codeQLQueryHistory.viewDil",
@@ -643,6 +647,11 @@
           "group": "9_qlCommands",
           "when": "view == codeQLQueryHistory"
         },
+        {
+          "command": "codeQLQueryHistory.exportCsvResults",
+          "group": "9_qlCommands",
+          "when": "view == codeQLQueryHistory && viewItem != interpretedResultsItem"
+        },
         {
           "command": "codeQLQueryHistory.viewCsvResults",
           "group": "9_qlCommands",
@@ -800,6 +809,10 @@
           "command": "codeQLQueryHistory.showQueryText",
           "when": "false"
         },
+        {
+          "command": "codeQLQueryHistory.exportCsvResults",
+          "when": "false"
+        },
         {
           "command": "codeQLQueryHistory.viewCsvResults",
           "when": "false"

diff --git a/extensions/ql-vscode/src/contextual/locationFinder.ts b/extensions/ql-vscode/src/contextual/locationFinder.ts
@@ -12,7 +12,7 @@ import { ProgressCallback } from '../commandRunner';
 import { KeyType } from './keyType';
 import { qlpackOfDatabase, resolveQueries } from './queryResolver';
 
-const SELECT_QUERY_NAME = '#select';
+export const SELECT_QUERY_NAME = '#select';
 export const TEMPLATE_NAME = 'selectedSourceFile';
 
 export interface FullLocationLink extends vscode.LocationLink {

diff --git a/extensions/ql-vscode/src/query-history.ts b/extensions/ql-vscode/src/query-history.ts
@@ -304,6 +304,12 @@ export class QueryHistoryManager extends DisposableObject {
         this.handleShowQueryText.bind(this)
       )
     );
+    this.push(
+      commandRunner(
+        'codeQLQueryHistory.exportCsvResults',
+        this.handleExportCsvResults.bind(this)
+      )
+    );
     this.push(
       commandRunner(
         'codeQLQueryHistory.viewCsvResults',
@@ -571,6 +577,32 @@ export class QueryHistoryManager extends DisposableObject {
     }
   }
 
+  async handleExportCsvResults(
+    singleItem: CompletedQuery,
+    multiSelect: CompletedQuery[]
+  ) {
+    if (!this.assertSingleQuery(multiSelect)) {
+      return;
+    }
+
+    const saveLocation = await vscode.window.showSaveDialog({
+      title: 'CSV Results',
+      saveLabel: 'Export',
+      filters: {
+        'Comma-separated values': ['csv'],
+      }
+    });
+    if (!saveLocation) {
+      void showAndLogErrorMessage('No save location selected for CSV export!');
+      return;
+    }
+    await singleItem.query.exportCsvResults(this.qs, saveLocation.fsPath, () => {
+      void this.tryOpenExternalFile(
+        saveLocation.fsPath
+      );
+    });
+  }
+
   async handleViewCsvResults(
     singleItem: CompletedQuery,
     multiSelect: CompletedQuery[]

diff --git a/extensions/ql-vscode/src/run-queries.ts b/extensions/ql-vscode/src/run-queries.ts
@@ -25,6 +25,8 @@ import * as qsClient from './queryserver-client';
 import { isQuickQueryPath } from './quick-query';
 import { compileDatabaseUpgradeSequence, hasNondestructiveUpgradeCapabilities, upgradeDatabaseExplicit } from './upgrades';
 import { ensureMetadataIsComplete } from './query-results';
+import { SELECT_QUERY_NAME } from './contextual/locationFinder';
+import { DecodedBqrsChunk } from './pure/bqrs-cli-types';
 
 /**
  * run-queries.ts
@@ -216,6 +218,29 @@ export class QueryInfo {
     return this.dilPath;
   }
 
+  async exportCsvResults(qs: qsClient.QueryServerClient, csvPath: string, onFinish: () => void): Promise<void> {
+    let stopDecoding = false;
+    const out = fs.createWriteStream(csvPath);
+    out.on('finish', onFinish);
+    out.on('error', () => {
+      if (!stopDecoding) {
+        stopDecoding = true;
+        void showAndLogErrorMessage(`Failed to write CSV results to ${csvPath}`);
+      }
+    });
+    let nextOffset: number | undefined = 0;
+    while (nextOffset !== undefined && !stopDecoding) {
+      const chunk: DecodedBqrsChunk = await qs.cliServer.bqrsDecode(this.resultsPaths.resultsPath, SELECT_QUERY_NAME, {
+        pageSize: 100,
+        offset: nextOffset,
+      });
+      for (const tuple of chunk.tuples)
+        out.write(tuple.join(',') + '\n');
+      nextOffset = chunk.next;
+    }
+    out.end();
+  }
+
   async ensureCsvProduced(qs: qsClient.QueryServerClient): Promise<string> {
     if (await this.hasCsv()) {
       return this.csvPath;