Add the commandRunner

extensions/ql-vscode/CHANGELOG.md

@@ -4,7 +4,9 @@
 
 - Add friendly welcome message when the databases view is empty.
 - Add open query, open results, and remove query commands in the query history view title bar.
-- Max number of simultaneous queries launchable by runQueries command is now configurable by changing the `codeQL.runningQueries.maxQueries` setting.
+- The maximum number of simultaneous queries launchable by the `CodeQL: Run Queries in Selected Files` command is now configurable by changing the `codeQL.runningQueries.maxQueries` setting.
+- Allow simultaneously run queries to be canceled in a single-click.
+- Prevent multiple upgrade dialogs from appearing when running simultaneous queries on upgradeable databases.
 - Fix sorting of results. Some pages of results would have the wrong sort order and columns.
 - Remember previous sort order when reloading query results.
 - Fix proper escaping of backslashes in SARIF message strings.

extensions/ql-vscode/src/astViewer.ts

@@ -3,7 +3,6 @@ import {
   ExtensionContext,
   TreeDataProvider,
   EventEmitter,
-  commands,
   Event,
   ProviderResult,
   TreeItemCollapsibleState,
@@ -19,7 +18,7 @@ import { DatabaseItem } from './databases';
 import { UrlValue, BqrsId } from './bqrs-cli-types';
 import { showLocation } from './interface-utils';
 import { isStringLoc, isWholeFileLoc, isLineColumnLoc } from './bqrs-utils';
-
+import { commandRunner } from './helpers';
 
 export interface AstItem {
   id: BqrsId;
@@ -45,10 +44,10 @@ class AstViewerDataProvider implements TreeDataProvider<AstItem> {
     this._onDidChangeTreeData.event;
 
   constructor() {
-    commands.registerCommand('codeQLAstViewer.gotoCode',
-      async (item: AstItem) => {
-        await showLocation(item.fileLocation);
-      });
+    commandRunner('codeQLAstViewer.gotoCode',
+    async (item: AstItem) => {
+      await showLocation(item.fileLocation);
+    });
   }
 
   refresh(): void {
@@ -109,7 +108,7 @@ export class AstViewer {
       showCollapseAll: true
     });
 
-    commands.registerCommand('codeQLAstViewer.clear', () => {
+    commandRunner('codeQLAstViewer.clear', async () => {
       this.clear();
     });
 

extensions/ql-vscode/src/contextual/locationFinder.ts

@@ -1,13 +1,14 @@
 import * as vscode from 'vscode';
 
 import { decodeSourceArchiveUri, zipArchiveScheme } from '../archive-filesystem-provider';
-import { ColumnKindCode, EntityValue, getResultSetSchema } from '../bqrs-cli-types';
+import { ColumnKindCode, EntityValue, getResultSetSchema, ResultSetSchema } from '../bqrs-cli-types';
 import { CodeQLCliServer } from '../cli';
 import { DatabaseManager, DatabaseItem } from '../databases';
 import fileRangeFromURI from './fileRangeFromURI';
 import * as messages from '../messages';
 import { QueryServerClient } from '../queryserver-client';
 import { QueryWithResults, compileAndRunQueryAgainstDatabase } from '../run-queries';
+import { ProgressCallback } from '../helpers';
 import { KeyType } from './keyType';
 import { qlpackOfDatabase, resolveQueries } from './queryResolver';
 
@@ -28,6 +29,8 @@ export interface FullLocationLink extends vscode.LocationLink {
  * @param dbm The database manager
  * @param uriString The selected source file and location
  * @param keyType The contextual query type to run
+ * @param progress A progress callback
+ * @param token A CancellationToken
  * @param filter A function that will filter extraneous results
  */
 export async function getLocationsForUriString(
@@ -36,37 +39,42 @@ export async function getLocationsForUriString(
   dbm: DatabaseManager,
   uriString: string,
   keyType: KeyType,
+  progress: ProgressCallback,
+  token: vscode.CancellationToken,
   filter: (src: string, dest: string) => boolean
 ): Promise<FullLocationLink[]> {
   const uri = decodeSourceArchiveUri(vscode.Uri.parse(uriString));
   const sourceArchiveUri = vscode.Uri.file(uri.sourceArchiveZipPath).with({ scheme: zipArchiveScheme });
 
   const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
-  if (db) {
-    const qlpack = await qlpackOfDatabase(cli, db);
-    if (qlpack === undefined) {
-      throw new Error('Can\'t infer qlpack from database source archive');
-    }
-    const links: FullLocationLink[] = [];
-    for (const query of await resolveQueries(cli, qlpack, keyType)) {
-      const templates: messages.TemplateDefinitions = {
-        [TEMPLATE_NAME]: {
-          values: {
-            tuples: [[{
-              stringValue: uri.pathWithinSourceArchive
-            }]]
-          }
-        }
-      };
-      const results = await compileAndRunQueryAgainstDatabase(cli, qs, db, false, vscode.Uri.file(query), templates);
-      if (results.result.resultType == messages.QueryResultType.SUCCESS) {
-        links.push(...await getLinksFromResults(results, cli, db, filter));
-      }
-    }
-    return links;
-  } else {
+  if (!db) {
     return [];
   }
+
+  const qlpack = await qlpackOfDatabase(cli, db);
+  if (qlpack === undefined) {
+    throw new Error('Can\'t infer qlpack from database source archive');
+  }
+  const templates = createTemplates(uri.pathWithinSourceArchive);
+
+  const links: FullLocationLink[] = [];
+  for (const query of await resolveQueries(cli, qlpack, keyType)) {
+    const results = await compileAndRunQueryAgainstDatabase(
+      cli,
+      qs,
+      db,
+      false,
+      vscode.Uri.file(query),
+      progress,
+      token,
+      templates
+    );
+
+    if (results.result.resultType == messages.QueryResultType.SUCCESS) {
+      links.push(...await getLinksFromResults(results, cli, db, filter));
+    }
+  }
+  return links;
 }
 
 async function getLinksFromResults(
@@ -79,10 +87,7 @@ async function getLinksFromResults(
   const bqrsPath = results.query.resultsPaths.resultsPath;
   const info = await cli.bqrsInfo(bqrsPath);
   const selectInfo = getResultSetSchema(SELECT_QUERY_NAME, info);
-  if (selectInfo && selectInfo.columns.length == 3
-    && selectInfo.columns[0].kind == ColumnKindCode.ENTITY
-    && selectInfo.columns[1].kind == ColumnKindCode.ENTITY
-    && selectInfo.columns[2].kind == ColumnKindCode.STRING) {
+  if (isValidSelect(selectInfo)) {
     // TODO: Page this
     const allTuples = await cli.bqrsDecode(bqrsPath, SELECT_QUERY_NAME);
     for (const tuple of allTuples.tuples) {
@@ -101,3 +106,22 @@ async function getLinksFromResults(
   }
   return localLinks;
 }
+
+function createTemplates(path: string): messages.TemplateDefinitions {
+  return {
+    [TEMPLATE_NAME]: {
+      values: {
+        tuples: [[{
+          stringValue: path
+        }]]
+      }
+    }
+  };
+}
+
+function isValidSelect(selectInfo: ResultSetSchema | undefined) {
+  return selectInfo && selectInfo.columns.length == 3
+    && selectInfo.columns[0].kind == ColumnKindCode.ENTITY
+    && selectInfo.columns[1].kind == ColumnKindCode.ENTITY
+    && selectInfo.columns[2].kind == ColumnKindCode.STRING;
+}

extensions/ql-vscode/src/contextual/templateProvider.ts

@@ -3,7 +3,7 @@ import * as vscode from 'vscode';
 import { decodeSourceArchiveUri, zipArchiveScheme } from '../archive-filesystem-provider';
 import { CodeQLCliServer } from '../cli';
 import { DatabaseManager } from '../databases';
-import { CachedOperation } from '../helpers';
+import { CachedOperation, ProgressCallback, withProgress } from '../helpers';
 import * as messages from '../messages';
 import { QueryServerClient } from '../queryserver-client';
 import { compileAndRunQueryAgainstDatabase, QueryWithResults } from '../run-queries';
@@ -44,14 +44,22 @@ export class TemplateQueryDefinitionProvider implements vscode.DefinitionProvide
   }
 
   private async getDefinitions(uriString: string): Promise<vscode.LocationLink[]> {
-    return getLocationsForUriString(
-      this.cli,
-      this.qs,
-      this.dbm,
-      uriString,
-      KeyType.DefinitionQuery,
-      (src, _dest) => src === uriString
-    );
+    return withProgress({
+      location: vscode.ProgressLocation.Notification,
+      cancellable: true,
+      title: 'Finding definitions'
+    }, async (progress, token) => {
+      return getLocationsForUriString(
+        this.cli,
+        this.qs,
+        this.dbm,
+        uriString,
+        KeyType.DefinitionQuery,
+        progress,
+        token,
+        (src, _dest) => src === uriString
+      );
+    });
   }
 }
 
@@ -83,14 +91,22 @@ export class TemplateQueryReferenceProvider implements vscode.ReferenceProvider
   }
 
   private async getReferences(uriString: string): Promise<FullLocationLink[]> {
-    return getLocationsForUriString(
-      this.cli,
-      this.qs,
-      this.dbm,
-      uriString,
-      KeyType.ReferenceQuery,
-      (_src, dest) => dest === uriString
-    );
+    return withProgress({
+      location: vscode.ProgressLocation.Notification,
+      cancellable: true,
+      title: 'Finding references'
+    }, async (progress, token) => {
+      return getLocationsForUriString(
+        this.cli,
+        this.qs,
+        this.dbm,
+        uriString,
+        KeyType.DefinitionQuery,
+        progress,
+        token,
+        (src, _dest) => src === uriString
+      );
+    });
   }
 }
 
@@ -101,6 +117,10 @@ export class TemplatePrintAstProvider {
     private cli: CodeQLCliServer,
     private qs: QueryServerClient,
     private dbm: DatabaseManager,
+
+    // Note: progress and token are only used if a cached value is not available
+    private progress: ProgressCallback,
+    private token: vscode.CancellationToken
   ) {
     this.cache = new CachedOperation<QueryWithResults | undefined>(this.getAst.bind(this));
   }
@@ -157,12 +177,15 @@ export class TemplatePrintAstProvider {
         }
       }
     };
+
     return await compileAndRunQueryAgainstDatabase(
       this.cli,
       this.qs,
       db,
       false,
       vscode.Uri.file(query),
+      this.progress,
+      this.token,
       templates
     );
   }