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/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## [UNRELEASED]
 
+- Fix bug where databases are not reregistered when the query server restarts. [#734](https://github.com/github/vscode-codeql/pull/734)
+
 ## 1.3.10 - 20 January 2021
 
 - Include the full stack in error log messages to help with debugging. [#726](https://github.com/github/vscode-codeql/pull/726)

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.reregisterDatabases.bind(this));
+
+    // Let this run async.
+    this.loadPersistedState();
   }
 
   public async openDatabase(
@@ -542,6 +545,22 @@ export class DatabaseManager extends DisposableObject {
     return databaseItem;
   }
 
+  private async reregisterDatabases(
+    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'
+      });
+    }));
+  }
+
   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

@@ -594,28 +594,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(
+    commandRunnerWithProgress('codeQL.chooseDatabaseFolder', (
+      progress: ProgressCallback,
+      token: CancellationToken
+    ) =>
+      databaseUI.handleChooseDatabaseFolder(progress, token), {
+      title: 'Choose a Database from a Folder'
     })
   );
   ctx.subscriptions.push(
-    commandRunner('codeQL.chooseDatabaseFolder', (
+    commandRunnerWithProgress('codeQL.chooseDatabaseArchive', (
       progress: ProgressCallback,
       token: CancellationToken
     ) =>
-      databaseUI.handleChooseDatabaseFolder(progress, token)
-    )
-  );
-  ctx.subscriptions.push(
-    commandRunner('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 queryServerStartListeners = [] 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.queryServerStartListeners.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.queryServerStartListeners.map(handler => handler(
+      progress,
+      token
+    )));
   }
 
   showLog(): void {

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

@@ -552,7 +552,7 @@ export async function compileAndRunQueryAgainstDatabase(
 
   const query = new QueryInfo(qlProgram, db, packConfig.dbscheme, quickEvalPosition, metadata, templates);
 
-  const upgradeDir = await tmp.dir({ dir: upgradesTmpDir.name });
+  const upgradeDir = await tmp.dir({ dir: upgradesTmpDir.name, unsafeCleanup: true });
   try {
     let upgradeQlo;
     if (await hasNondestructiveUpgradeCapabilities(qs)) {
@@ -615,7 +615,11 @@ export async function compileAndRunQueryAgainstDatabase(
       return createSyntheticResult(query, db, historyItemOptions, 'Query had compilation errors', messages.QueryResultType.OTHER_ERROR);
     }
   } finally {
-    upgradeDir.cleanup();
+    try {
+      upgradeDir.cleanup();
+    } catch (e) {
+      qs.logger.log(`Could not clean up the upgrades dir. Reason: ${e.message || e}`);
+    }
   }
 }
 

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';
@@ -14,6 +14,7 @@ import { compileAndRunQueryAgainstDatabase } from '../../run-queries';
 import { CodeQLCliServer } from '../../cli';
 import { QueryServerClient } from '../../queryserver-client';
 import { skipIfNoCodeQL } from '../ensureCli';
+import { QueryResultType } from '../../pure/messages';
 
 
 /**
@@ -94,10 +95,35 @@ describe('Queries', function() {
       // just check that the query was successful
       expect(result.database.name).to.eq('db');
       expect(result.options.queryText).to.eq(fs.readFileSync(queryPath, 'utf8'));
+      expect(result.result.resultType).to.eq(QueryResultType.SUCCESS);
     } catch (e) {
       console.error('Test Failed');
       fail(e);
     }
   });
 
+  // Asserts a fix for bug https://github.com/github/vscode-codeql/issues/733
+  it('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'));
+      expect(result.result.resultType).to.eq(QueryResultType.SUCCESS);
+    } catch (e) {
+      console.error('Test Failed');
+      fail(e);
+    }
+  });
 });

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

@@ -33,7 +33,10 @@ class Checkpoint<T> {
   constructor() {
     this.res = () => { /**/ };
     this.rej = () => { /**/ };
-    this.promise = new Promise((res, rej) => { this.res = res; this.rej = rej; });
+    this.promise = new Promise((res, rej) => {
+      this.res = res as () => {};
+      this.rej = rej;
+    });
   }
 
   async done(): Promise<T> {
@@ -81,6 +84,11 @@ const queryTestCases: QueryTestCase[] = [
   }
 ];
 
+const db: messages.Dataset = {
+  dbDir: path.join(__dirname, '../test-db'),
+  workingSet: 'default',
+};
+
 describe('using the query server', function() {
   before(function() {
     skipIfNoCodeQL(this);
@@ -120,6 +128,12 @@ describe('using the query server', function() {
     const evaluationSucceeded = new Checkpoint<void>();
     const parsedResults = new Checkpoint<void>();
 
+    it('should register the database if necessary', async () => {
+      if (await qs.supportsDatabaseRegistration()) {
+        await qs.sendRequest(messages.registerDatabases, { databases: [db] }, token, (() => { /**/ }) as any);
+      }
+    });
+
     it(`should be able to compile query ${queryName}`, async function() {
       await queryServerStarted.done();
       expect(fs.existsSync(queryTestCase.queryPath)).to.be.true;
@@ -166,15 +180,11 @@ describe('using the query server', function() {
           id: callbackId,
           timeoutSecs: 1000,
         };
-        const db: messages.Dataset = {
-          dbDir: path.join(__dirname, '../test-db'),
-          workingSet: 'default',
-        };
         const params: messages.EvaluateQueriesParams = {
           db,
           evaluateId: callbackId,
           queries: [queryToRun],
-          stopOnError: false,
+          stopOnError: true,
           useSequenceHint: false
         };
         await qs.sendRequest(messages.runQueries, params, token, () => { /**/ });

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

@@ -28,7 +28,10 @@ import { workspace } from 'vscode';
 process.on('unhandledRejection', e => {
   console.error('Unhandled rejection.');
   console.error(e);
-  process.exit(-1);
+  // Must use a setTimeout in order to ensure the log is fully flushed before exiting
+  setTimeout(() => {
+    process.exit(-1);
+  }, 2000);
 });
 
 const _1MB = 1024 * 1024;
@@ -36,7 +39,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

extensions/ql-vscode/src/vscode-tests/minimal-workspace/databases.test.ts

@@ -66,7 +66,8 @@ describe('databases', () => {
       } as unknown as ExtensionContext,
       {
         sendRequest: sendRequestSpy,
-        supportsDatabaseRegistration: supportsDatabaseRegistrationSpy
+        supportsDatabaseRegistration: supportsDatabaseRegistrationSpy,
+        onDidStartQueryServer: () => { /**/ }
       } as unknown as QueryServerClient,
       {
         supportsLanguageName: supportsLanguageNameSpy,