Ensure databases are re-registered when query server restarts

This commit fixes #733. It does it by ensuring that the query server
emits an event when it restarts the query server. The database manager
listens for this even and properly re-registers its databases.

A few caveats though:

1. Convert query restarts to using a command that includes progress.
   This will ensure that errors on restart are logged properly.
2. Because we want to log errors, we cannot use the vscode standard
   EventEmitters. They run in the next tick and therefore any errors
   will not be associated with this command execution.
3. Update the default cli version to run integration tests against to
   2.4.2.
4. Add a new integration test that fails if databases are not
   re-registered.

.github/workflows/main.yml

@@ -121,7 +121,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, windows-latest]
-        version: ['v2.2.6', 'v2.3.3', 'v2.4.0']
+        version: ['v2.2.6', 'v2.3.3', 'v2.4.2']
     env:
       CLI_VERSION: ${{ matrix.version }}
       TEST_CODEQL_PATH: '${{ github.workspace }}/codeql'

extensions/ql-vscode/src/databases.ts

@@ -514,7 +514,10 @@ export class DatabaseManager extends DisposableObject {
   ) {
     super();
 
-    this.loadPersistedState();  // Let this run async.
+    qs.onDidStartQueryServer(this.onQueryServerRestarted);
+
+    // Let this run async.
+    this.loadPersistedState();
   }
 
   public async openDatabase(
@@ -542,6 +545,24 @@ export class DatabaseManager extends DisposableObject {
     return databaseItem;
   }
 
+  private onQueryServerRestarted = async (
+    progress: ProgressCallback,
+    token: vscode.CancellationToken
+  ) => {
+    let completed = 0;
+    await Promise.all(this._databaseItems.map(async (databaseItem) => {
+      await this.registerDatabase(progress, token, databaseItem);
+      completed++;
+      progress({
+        maxStep: this._databaseItems.length,
+        step: completed,
+        message: 'Re-registering databases'
+      });
+    }));
+
+    throw new Error('yikes');
+  }
+
   private async addDatabaseSourceArchiveFolder(item: DatabaseItem) {
     // The folder may already be in workspace state from a previous
     // session. If not, add it.

extensions/ql-vscode/src/extension.ts

@@ -595,28 +595,36 @@ async function activateWithInstalledDistribution(
   );
 
   ctx.subscriptions.push(
-    commandRunner('codeQL.restartQueryServer', async () => {
-      await qs.restartQueryServer();
+    commandRunnerWithProgress('codeQL.restartQueryServer', async (
+      progress: ProgressCallback,
+      token: CancellationToken
+    ) => {
+      await qs.restartQueryServer(progress, token);
       helpers.showAndLogInformationMessage('CodeQL Query Server restarted.', {
         outputLogger: queryServerLogger,
       });
+    }, {
+      title: 'Restarting Query Server'
     })
   );
+
   ctx.subscriptions.push(
-    commandRunner('codeQL.chooseDatabaseFolder', (
+    commandRunnerWithProgress('codeQL.chooseDatabaseFolder', (
       progress: ProgressCallback,
       token: CancellationToken
     ) =>
-      databaseUI.handleChooseDatabaseFolder(progress, token)
-    )
+      databaseUI.handleChooseDatabaseFolder(progress, token), {
+      title: 'Choose a Database from a Folder'
+    })
   );
   ctx.subscriptions.push(
-    commandRunner('codeQL.chooseDatabaseArchive', (
+    commandRunnerWithProgress('codeQL.chooseDatabaseArchive', (
       progress: ProgressCallback,
       token: CancellationToken
     ) =>
-      databaseUI.handleChooseDatabaseArchive(progress, token)
-    )
+      databaseUI.handleChooseDatabaseArchive(progress, token), {
+      title: 'Choose a Database from an Archive'
+    })
   );
   ctx.subscriptions.push(
     commandRunnerWithProgress('codeQL.chooseDatabaseLgtm', (

extensions/ql-vscode/src/queryserver-client.ts

@@ -1,14 +1,15 @@
 import * as cp from 'child_process';
 import * as path from 'path';
 import { DisposableObject } from './vscode-utils/disposable-object';
-import { Disposable } from 'vscode';
-import { CancellationToken, createMessageConnection, MessageConnection, RequestType } from 'vscode-jsonrpc';
+import { Disposable, CancellationToken, commands } from 'vscode';
+import { createMessageConnection, MessageConnection, RequestType } from 'vscode-jsonrpc';
 import * as cli from './cli';
 import { QueryServerConfig } from './config';
 import { Logger, ProgressReporter } from './logging';
 import { completeQuery, EvaluationResult, progress, ProgressMessage, WithProgressId } from './pure/messages';
 import * as messages from './pure/messages';
 import { SemVer } from 'semver';
+import { ProgressCallback, ProgressTask } from './commandRunner';
 
 type ServerOpts = {
   logger: Logger;
@@ -60,6 +61,16 @@ export class QueryServerClient extends DisposableObject {
   nextCallback: number;
   nextProgress: number;
   withProgressReporting: WithProgressReporting;
+
+  private readonly _onDidStartQueryServer = [] as ProgressTask<void>[];
+
+  // Can't use standard vscode EventEmitter here since they do not cause the calling
+  // function to fail if one of the event handlers fail. This is something that
+  // we need here.
+  readonly onDidStartQueryServer = (e: ProgressTask<void>) => {
+    this._onDidStartQueryServer.push(e);
+  }
+
   public activeQueryName: string | undefined;
 
   constructor(
@@ -71,10 +82,8 @@ export class QueryServerClient extends DisposableObject {
     super();
     // When the query server configuration changes, restart the query server.
     if (config.onDidChangeConfiguration !== undefined) {
-      this.push(config.onDidChangeConfiguration(async () => {
-        this.logger.log('Restarting query server due to configuration changes...');
-        await this.restartQueryServer();
-      }, this));
+      this.push(config.onDidChangeConfiguration(() =>
+        commands.executeCommand('codeQL.restartQueryServer')));
     }
     this.withProgressReporting = withProgressReporting;
     this.nextCallback = 0;
@@ -97,9 +106,19 @@ export class QueryServerClient extends DisposableObject {
   }
 
   /** Restarts the query server by disposing of the current server process and then starting a new one. */
-  async restartQueryServer(): Promise<void> {
+  async restartQueryServer(
+    progress: ProgressCallback,
+    token: CancellationToken
+  ): Promise<void> {
     this.stopQueryServer();
     await this.startQueryServer();
+
+    // Ensure we await all responses from event handlers so that
+    // errors can be properly reported to the user.
+    await Promise.all(this._onDidStartQueryServer.map(handler => handler(
+      progress,
+      token
+    )));
   }
 
   showLog(): void {

extensions/ql-vscode/src/vscode-tests/cli-integration/queries.test.ts

@@ -1,5 +1,5 @@
 import { fail } from 'assert';
-import { CancellationToken, extensions, Uri } from 'vscode';
+import { CancellationToken, commands, extensions, Uri } from 'vscode';
 import * as sinon from 'sinon';
 import * as path from 'path';
 import * as fs from 'fs-extra';
@@ -100,4 +100,27 @@ describe('Queries', function() {
     }
   });
 
+  // Asserts a fix for bug https://github.com/github/vscode-codeql/issues/733
+  it.only('should restart the database and run a query', async () => {
+    try {
+      await commands.executeCommand('codeQL.restartQueryServer');
+      const queryPath = path.join(__dirname, 'data', 'simple-query.ql');
+      const result = await compileAndRunQueryAgainstDatabase(
+        cli,
+        qs,
+        dbItem,
+        false,
+        Uri.file(queryPath),
+        progress,
+        token
+      );
+
+      // this message would indicate that the databases were not properly reregistered
+      expect(result.result.message).not.to.eq('No result from server');
+      expect(result.options.queryText).to.eq(fs.readFileSync(queryPath, 'utf8'));
+    } catch (e) {
+      console.error('Test Failed');
+      fail(e);
+    }
+  });
 });

extensions/ql-vscode/src/vscode-tests/ensureCli.ts

@@ -36,7 +36,7 @@ const _10MB = _1MB * 10;
 
 // CLI version to test. Hard code the latest as default. And be sure
 // to update the env if it is not otherwise set.
-const CLI_VERSION = process.env.CLI_VERSION || 'v2.4.0';
+const CLI_VERSION = process.env.CLI_VERSION || 'v2.4.2';
 process.env.CLI_VERSION = CLI_VERSION;
 
 // Base dir where CLIs will be downloaded into