Add the commandRunner

The commandRunner wraps all vscode command registrations. It provides
uniform error handling and an optional progress monitor.

In general, progress monitors should only be created by the
commandRunner and passed through to the locations that use it.

Author: aeisenberg

extensions/ql-vscode/src/astViewer.ts

@@ -3,7 +3,6 @@ import {
   ExtensionContext,
   TreeDataProvider,
   EventEmitter,
-  commands,
   Event,
   ProviderResult,
   TreeItemCollapsibleState,
@@ -19,6 +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 {
@@ -45,10 +45,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 +109,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 await 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 await 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
     );
   }