Ensure databases are re-registered when query server restarts

.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,9 @@
 
 ## [UNRELEASED]
 
+- Fix bug where databases are not reregistered when the query server restarts. [#734](https://github.com/github/vscode-codeql/pull/734)
+- Fix bug where upgrade requests were erroneously being marked as failed. [#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(
-    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/pure/messages.ts

@@ -479,7 +479,7 @@ export interface CompileUpgradeSequenceResult {
   /**
    * The compiled upgrades as a single file.
    */
-  compiledUpgrades?: string;
+  compiledUpgrade?: string;
   /**
    * Any errors that occurred when checking the scripts.
    */

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

@@ -354,14 +354,14 @@ async function compileNonDestructiveUpgrade(
     reportNoUpgradePath(query);
   }
   const result = await compileDatabaseUpgradeSequence(qs, query.dbItem, scripts, upgradeTemp, progress, token);
-  if (result.compiledUpgrades === undefined) {
+  if (result.compiledUpgrade === undefined) {
     const error = result.error || '[no error message available]';
     throw new Error(error);
   }
   // We can upgrade to the actual target
   query.program.dbschemePath = query.queryDbscheme;
   // We are new enough that we will always support single file upgrades.
-  return result.compiledUpgrades;
+  return result.compiledUpgrade;
 
 }
 
@@ -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/upgrades.ts

@@ -20,9 +20,9 @@ const MAX_UPGRADE_MESSAGE_LINES = 10;
 
 /**
  * Check that we support non-destructive upgrades.
- * 
- * This requires 3 features. The ability to compile an upgrade sequence; The ability to 
- * run a non-destructive upgrades as a query; the ability to specify a target when 
+ *
+ * This requires 3 features. The ability to compile an upgrade sequence; The ability to
+ * run a non-destructive upgrades as a query; the ability to specify a target when
  * resolving upgrades. We check for a version of codeql that has all three features.
  */
 export async function hasNondestructiveUpgradeCapabilities(qs: qsClient.QueryServerClient): Promise<boolean> {
@@ -34,12 +34,14 @@ export async function hasNondestructiveUpgradeCapabilities(qs: qsClient.QuerySer
  * Compile a database upgrade sequence.
  * Callers must check that this is valid with the current queryserver first.
  */
-export async function compileDatabaseUpgradeSequence(qs: qsClient.QueryServerClient,
+export async function compileDatabaseUpgradeSequence(
+  qs: qsClient.QueryServerClient,
   db: DatabaseItem,
   resolvedSequence: string[],
   currentUpgradeTmp: tmp.DirectoryResult,
   progress: ProgressCallback,
-  token: vscode.CancellationToken): Promise<messages.CompileUpgradeSequenceResult> {
+  token: vscode.CancellationToken
+): Promise<messages.CompileUpgradeSequenceResult> {
   if (db.contents === undefined || db.contents.dbSchemeUri === undefined) {
     throw new Error('Database is invalid, and cannot be upgraded.');
   }

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,15 +28,18 @@ 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;
 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,