Ensure database upgrade request happens only once

When a user runs multiple queries on a non-upgraded database, ensure
that only one dialog appears for upgrade.

This commit also migrates the upgrades.ts file to using the passed-in
cancellation token and progress monitor. This ensures that cancelling
a database upgrade command will also cancel out of any wrapper
operations.

Fixes #534

Author: aeisenberg

extensions/ql-vscode/CHANGELOG.md

@@ -5,6 +5,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.
+- Allow simultaneously run queries to be canceled in a single-click.
+- Prevent multiple upgrade dialogs from appearing when running simultaneous queries on upgradeable databases.
+- Max number of simultaneous queries launchable by runQueries command is now configurable by changing the codeQL.runningQueries.maxQueries setting.
 - 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/databases-ui.ts

@@ -246,7 +246,12 @@ export class DatabaseUI extends DisposableObject {
     ctx.subscriptions.push(
       commandRunner(
         'codeQL.upgradeCurrentDatabase',
-        this.handleUpgradeCurrentDatabase
+        this.handleUpgradeCurrentDatabase,
+        {
+          location: ProgressLocation.Notification,
+          title: 'Upgrading current database',
+          cancellable: true,
+        }
       )
     );
     ctx.subscriptions.push(
@@ -263,13 +268,23 @@ export class DatabaseUI extends DisposableObject {
     ctx.subscriptions.push(
       commandRunner(
         'codeQLDatabases.chooseDatabaseFolder',
-        this.handleChooseDatabaseFolder
+        this.handleChooseDatabaseFolder,
+        {
+          location: ProgressLocation.Notification,
+          title: 'Adding database from folder',
+          cancellable: false,
+        }
       )
     );
     ctx.subscriptions.push(
       commandRunner(
         'codeQLDatabases.chooseDatabaseArchive',
-        this.handleChooseDatabaseArchive
+        this.handleChooseDatabaseArchive,
+        {
+          location: ProgressLocation.Notification,
+          title: 'Adding database from archive',
+          cancellable: false,
+        }
       )
     );
     ctx.subscriptions.push(
@@ -320,7 +335,12 @@ export class DatabaseUI extends DisposableObject {
     ctx.subscriptions.push(
       commandRunner(
         'codeQLDatabases.upgradeDatabase',
-        this.handleUpgradeDatabase
+        this.handleUpgradeDatabase,
+        {
+          location: ProgressLocation.Notification,
+          title: 'Upgrading database',
+          cancellable: true,
+        }
       )
     );
     ctx.subscriptions.push(
@@ -393,6 +413,13 @@ export class DatabaseUI extends DisposableObject {
     );
   };
 
+  async tryUpgradeCurrentDatabase(
+    progress: ProgressCallback,
+    token: CancellationToken
+  ) {
+    await this.handleUpgradeCurrentDatabase(progress, token);
+  }
+
   private handleSortByName = async () => {
     if (this.treeDataProvider.sortOrder === SortOrder.NameAsc) {
       this.treeDataProvider.sortOrder = SortOrder.NameDesc;
@@ -409,45 +436,47 @@ export class DatabaseUI extends DisposableObject {
     }
   };
 
-  private handleUpgradeCurrentDatabase = async (): Promise<void> => {
+  private handleUpgradeCurrentDatabase = async (
+    progress: ProgressCallback,
+    token: CancellationToken,
+  ): Promise<void> => {
     await this.handleUpgradeDatabase(
+      progress, token,
       this.databaseManager.currentDatabaseItem,
       []
     );
   };
 
   private handleUpgradeDatabase = async (
+    progress: ProgressCallback,
+    token: CancellationToken,
     databaseItem: DatabaseItem | undefined,
-    multiSelect: DatabaseItem[] | undefined
+    multiSelect: DatabaseItem[] | undefined,
   ): Promise<void> => {
     if (multiSelect?.length) {
       await Promise.all(
-        multiSelect.map((dbItem) => this.handleUpgradeDatabase(dbItem, []))
+        multiSelect.map((dbItem) => this.handleUpgradeDatabase(progress, token, dbItem, []))
       );
     }
     if (this.queryServer === undefined) {
-      logger.log(
+      throw new Error(
         'Received request to upgrade database, but there is no running query server.'
       );
-      return;
     }
     if (databaseItem === undefined) {
-      logger.log(
+      throw new Error(
         'Received request to upgrade database, but no database was provided.'
       );
-      return;
     }
     if (databaseItem.contents === undefined) {
-      logger.log(
+      throw new Error(
         'Received request to upgrade database, but database contents could not be found.'
       );
-      return;
     }
     if (databaseItem.contents.dbSchemeUri === undefined) {
-      logger.log(
+      throw new Error(
         'Received request to upgrade database, but database has no schema.'
       );
-      return;
     }
 
     // Search for upgrade scripts in any workspace folders available
@@ -461,16 +490,17 @@ export class DatabaseUI extends DisposableObject {
     const { scripts, finalDbscheme } = upgradeInfo;
 
     if (finalDbscheme === undefined) {
-      logger.log('Could not determine target dbscheme to upgrade to.');
-      return;
+      throw new Error('Could not determine target dbscheme to upgrade to.');
     }
     const targetDbSchemeUri = Uri.file(finalDbscheme);
 
     await upgradeDatabase(
       this.queryServer,
       databaseItem,
       targetDbSchemeUri,
-      getUpgradesDirectories(scripts)
+      getUpgradesDirectories(scripts),
+      progress,
+      token
     );
   };
 

extensions/ql-vscode/src/extension.ts

@@ -170,7 +170,7 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
             cancellable: false,
           };
 
-          // Avoid using commandRunner here because this function might be called upon extension activation
+          // Avoid using commandRunner here because this function is called upon extension activation
           await helpers.withProgress(progressOptions, progress =>
             distributionManager.installExtensionManagedDistributionRelease(result.updatedRelease, progress));
 
@@ -512,11 +512,21 @@ async function activateWithInstalledDistribution(
           });
         }
 
+        if (queryUris.length > 1) {
+          // Try to upgrade the current database before running any queries
+          // so that the user isn't confronted with multiple upgrade
+          // requests for each query to run.
+          // Only do it if running multiple queries since this check is
+          // performed on each query run anyway.
+          await databaseUI.tryUpgradeCurrentDatabase(progress, token);
+        }
+
         wrappedProgress({
           maxStep: queryUris.length,
           step: queryUris.length - queriesRemaining,
           message: ''
         });
+
         await Promise.all(queryUris.map(async uri =>
           compileAndRunQuery(false, uri, wrappedProgress, token)
             .then(() => queriesRemaining--)
@@ -550,6 +560,7 @@ async function activateWithInstalledDistribution(
       displayQuickQuery(ctx, cliServer, databaseUI, progress, token)
     )
   );
+
   ctx.subscriptions.push(
     helpers.commandRunner('codeQL.restartQueryServer', async () => {
       await qs.restartQueryServer();

extensions/ql-vscode/src/run-queries.ts

@@ -258,7 +258,9 @@ async function getSelectedPosition(editor: TextEditor): Promise<messages.Positio
 async function checkDbschemeCompatibility(
   cliServer: cli.CodeQLCliServer,
   qs: qsClient.QueryServerClient,
-  query: QueryInfo
+  query: QueryInfo,
+  progress: helpers.ProgressCallback,
+  token: CancellationToken,
 ): Promise<void> {
   const searchPath = helpers.getOnDiskWorkspaceFolders();
 
@@ -293,7 +295,9 @@ async function checkDbschemeCompatibility(
         qs,
         query.dbItem,
         Uri.file(finalDbscheme),
-        getUpgradesDirectories(scripts)
+        getUpgradesDirectories(scripts),
+        progress,
+        token
       );
     }
   }
@@ -481,7 +485,7 @@ export async function compileAndRunQueryAgainstDatabase(
   }
 
   const query = new QueryInfo(qlProgram, db, packConfig.dbscheme, quickEvalPosition, metadata, templates);
-  await checkDbschemeCompatibility(cliServer, qs, query);
+  await checkDbschemeCompatibility(cliServer, qs, query, progress, token);
 
   let errors;
   try {

extensions/ql-vscode/src/upgrades.ts

@@ -20,11 +20,15 @@ const MAX_UPGRADE_MESSAGE_LINES = 10;
  * @returns the `UpgradeParams` needed to start the upgrade, if the upgrade is possible and was confirmed by the user, or `undefined` otherwise.
  */
 async function checkAndConfirmDatabaseUpgrade(
-  qs: qsClient.QueryServerClient, db: DatabaseItem, targetDbScheme: vscode.Uri, upgradesDirectories: vscode.Uri[]
+  qs: qsClient.QueryServerClient,
+  db: DatabaseItem,
+  targetDbScheme: vscode.Uri,
+  upgradesDirectories: vscode.Uri[],
+  progress: helpers.ProgressCallback,
+  token: vscode.CancellationToken,
 ): Promise<messages.UpgradeParams | undefined> {
   if (db.contents === undefined || db.contents.dbSchemeUri === undefined) {
-    helpers.showAndLogErrorMessage('Database is invalid, and cannot be upgraded.');
-    return;
+    throw new Error('Database is invalid, and cannot be upgraded.');
   }
   const params: messages.UpgradeParams = {
     fromDbscheme: db.contents.dbSchemeUri.fsPath,
@@ -35,11 +39,10 @@ async function checkAndConfirmDatabaseUpgrade(
   let checkUpgradeResult: messages.CheckUpgradeResult;
   try {
     qs.logger.log('Checking database upgrade...');
-    checkUpgradeResult = await checkDatabaseUpgrade(qs, params);
+    checkUpgradeResult = await checkDatabaseUpgrade(qs, params, progress, token);
   }
   catch (e) {
-    helpers.showAndLogErrorMessage(`Database cannot be upgraded: ${e}`);
-    return;
+    throw new Error(`Database cannot be upgraded: ${e}`);
   }
   finally {
     qs.logger.log('Done checking database upgrade.');
@@ -48,12 +51,15 @@ async function checkAndConfirmDatabaseUpgrade(
   const checkedUpgrades = checkUpgradeResult.checkedUpgrades;
   if (checkedUpgrades === undefined) {
     const error = checkUpgradeResult.upgradeError || '[no error message available]';
-    await helpers.showAndLogErrorMessage(`Database cannot be upgraded: ${error}`);
-    return;
+    throw new Error(`Database cannot be upgraded: ${error}`);
   }
 
   if (checkedUpgrades.scripts.length === 0) {
-    await helpers.showAndLogInformationMessage('Database is already up to date; nothing to do.');
+    progress({
+      step: 3,
+      maxStep: 3,
+      message: 'Database is already up to date; nothing to do.'
+    });
     return;
   }
 
@@ -114,17 +120,19 @@ async function checkAndConfirmDatabaseUpgrade(
 export async function upgradeDatabase(
   qs: qsClient.QueryServerClient,
   db: DatabaseItem, targetDbScheme: vscode.Uri,
-  upgradesDirectories: vscode.Uri[]
+  upgradesDirectories: vscode.Uri[],
+  progress: helpers.ProgressCallback,
+  token: vscode.CancellationToken,
 ): Promise<messages.RunUpgradeResult | undefined> {
-  const upgradeParams = await checkAndConfirmDatabaseUpgrade(qs, db, targetDbScheme, upgradesDirectories);
+  const upgradeParams = await checkAndConfirmDatabaseUpgrade(qs, db, targetDbScheme, upgradesDirectories, progress, token);
 
   if (upgradeParams === undefined) {
     return;
   }
 
   let compileUpgradeResult: messages.CompileUpgradeResult;
   try {
-    compileUpgradeResult = await compileDatabaseUpgrade(qs, upgradeParams);
+    compileUpgradeResult = await compileDatabaseUpgrade(qs, upgradeParams, progress, token);
   }
   catch (e) {
     helpers.showAndLogErrorMessage(`Compilation of database upgrades failed: ${e}`);
@@ -143,7 +151,7 @@ export async function upgradeDatabase(
   try {
     qs.logger.log('Running the following database upgrade:');
     qs.logger.log(compileUpgradeResult.compiledUpgrades.scripts.map(s => s.description.description).join('\n'));
-    return await runDatabaseUpgrade(qs, db, compileUpgradeResult.compiledUpgrades);
+    return await runDatabaseUpgrade(qs, db, compileUpgradeResult.compiledUpgrades, progress, token);
   }
   catch (e) {
     helpers.showAndLogErrorMessage(`Database upgrade failed: ${e}`);
@@ -155,34 +163,46 @@ export async function upgradeDatabase(
 }
 
 async function checkDatabaseUpgrade(
-  qs: qsClient.QueryServerClient, upgradeParams: messages.UpgradeParams
+  qs: qsClient.QueryServerClient,
+  upgradeParams: messages.UpgradeParams,
+  progress: helpers.ProgressCallback,
+  token: vscode.CancellationToken,
 ): Promise<messages.CheckUpgradeResult> {
-  // Avoid using commandRunner here because this function might be called upon extension activation
-  return helpers.withProgress({
-    location: vscode.ProgressLocation.Notification,
-    title: 'Checking for database upgrades',
-    cancellable: true,
-  }, (progress, token) => qs.sendRequest(messages.checkUpgrade, upgradeParams, token, progress));
+  progress({
+    step: 1,
+    maxStep: 3,
+    message: 'Checking for database upgrades'
+  });
+
+  return qs.sendRequest(messages.checkUpgrade, upgradeParams, token, progress);
 }
 
 async function compileDatabaseUpgrade(
-  qs: qsClient.QueryServerClient, upgradeParams: messages.UpgradeParams
+  qs: qsClient.QueryServerClient,
+  upgradeParams: messages.UpgradeParams,
+  progress: helpers.ProgressCallback,
+  token: vscode.CancellationToken,
 ): Promise<messages.CompileUpgradeResult> {
   const params: messages.CompileUpgradeParams = {
     upgrade: upgradeParams,
     upgradeTempDir: upgradesTmpDir.name
   };
 
-  // Avoid using commandRunner here because this function might be called upon extension activation
-  return helpers.withProgress({
-    location: vscode.ProgressLocation.Notification,
-    title: 'Compiling database upgrades',
-    cancellable: true,
-  }, (progress, token) => qs.sendRequest(messages.compileUpgrade, params, token, progress));
+  progress({
+    step: 2,
+    maxStep: 3,
+    message: 'Compiling database upgrades'
+  });
+
+  return qs.sendRequest(messages.compileUpgrade, params, token, progress);
 }
 
 async function runDatabaseUpgrade(
-  qs: qsClient.QueryServerClient, db: DatabaseItem, upgrades: messages.CompiledUpgrades
+  qs: qsClient.QueryServerClient,
+  db: DatabaseItem,
+  upgrades: messages.CompiledUpgrades,
+  progress: helpers.ProgressCallback,
+  token: vscode.CancellationToken,
 ): Promise<messages.RunUpgradeResult> {
 
   if (db.contents === undefined || db.contents.datasetUri === undefined) {
@@ -199,10 +219,5 @@ async function runDatabaseUpgrade(
     toRun: upgrades
   };
 
-  // Avoid using commandRunner here because this function might be called upon extension activation
-  return helpers.withProgress({
-    location: vscode.ProgressLocation.Notification,
-    title: 'Running database upgrades',
-    cancellable: true,
-  }, (progress, token) => qs.sendRequest(messages.runUpgrade, params, token, progress));
+  return qs.sendRequest(messages.runUpgrade, params, token, progress);
 }