QueryServer: Add support for new query-server (#1508)

* QueryServer: Add support for new query-server

* Add a new canary flag to enable new query server support

* Add evaluation results to query object

Ensures better backwards compatibility with legacy query objects.

* Fix query server command name

* Add log message for new query server

* Use only legacy results

Co-authored-by: alexet <alexet@semmle.com>

extensions/ql-vscode/src/cli.ts

@@ -11,7 +11,7 @@ import { promisify } from 'util';
 import { CancellationToken, commands, Disposable, Uri } from 'vscode';
 
 import { BQRSInfo, DecodedBqrsChunk } from './pure/bqrs-cli-types';
-import { CliConfig } from './config';
+import { allowCanaryQueryServer, CliConfig } from './config';
 import { DistributionProvider, FindDistributionResultKind } from './distribution';
 import { assertNever, getErrorMessage, getErrorStack } from './pure/helpers-pure';
 import { QueryMetadata, SortDirection } from './pure/interface-types';
@@ -1330,6 +1330,11 @@ export class CliVersionConstraint {
    */
   public static CLI_VERSION_WITH_SOURCEMAP = new SemVer('2.10.3');
 
+  /**
+   * CLI version that supports the new query server.
+   */
+  public static CLI_VERSION_WITH_NEW_QUERY_SERVER = new SemVer('2.11.0');
+
   constructor(private readonly cli: CodeQLCliServer) {
     /**/
   }
@@ -1405,4 +1410,12 @@ export class CliVersionConstraint {
   async supportsSourceMap() {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_SOURCEMAP);
   }
+
+  async supportsNewQueryServer() {
+    // TODO while under development, users _must_ opt-in to the new query server
+    // by setting the `codeql.canaryQueryServer` setting to `true`.
+    // Ignore the version check for now.
+    return allowCanaryQueryServer();
+    // return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_NEW_QUERY_SERVER);
+  }
 }

extensions/ql-vscode/src/config.ts

@@ -317,6 +317,17 @@ export function isCanary() {
   return !!CANARY_FEATURES.getValue<boolean>();
 }
 
+/**
+ * Enables the experimental query server
+ */
+export const CANARY_QUERY_SERVER = new Setting('canaryQueryServer', ROOT_SETTING);
+
+
+export function allowCanaryQueryServer() {
+  return !!CANARY_QUERY_SERVER.getValue<boolean>();
+}
+
+
 /**
  * Avoids caching in the AST viewer if the user is also a canary user.
  */
@@ -343,12 +354,12 @@ export async function setRemoteRepositoryLists(lists: Record<string, string[]> |
 }
 
 /**
- * Path to a file that contains lists of GitHub repositories that you want to query remotely via 
+ * Path to a file that contains lists of GitHub repositories that you want to query remotely via
  * the "Run Variant Analysis" command.
  * Note: This command is only available for internal users.
- * 
+ *
  * This setting should be a path to a JSON file that contains a JSON object where each key is a
- * user-specified name (string), and the value is an array of GitHub repositories 
+ * user-specified name (string), and the value is an array of GitHub repositories
  * (of the form `<owner>/<repo>`).
  */
 const REPO_LISTS_PATH = new Setting('repositoryListsPath', REMOTE_QUERIES_SETTING);

extensions/ql-vscode/src/extension.ts

@@ -73,7 +73,8 @@ import { WebviewReveal } from './interface-utils';
 import { ideServerLogger, logger, ProgressReporter, queryServerLogger } from './logging';
 import { QueryHistoryManager } from './query-history';
 import { CompletedLocalQueryInfo, LocalQueryInfo } from './query-results';
-import * as qsClient from './legacy-query-server/queryserver-client';
+import * as legacyQueryServer from './legacy-query-server/queryserver-client';
+import * as newQueryServer from './query-server/queryserver-client';
 import { displayQuickQuery } from './quick-query';
 import { QLTestAdapterFactory } from './test-adapter';
 import { TestUIService } from './test-ui';
@@ -103,6 +104,7 @@ import { JoinOrderScannerProvider } from './log-insights/join-order';
 import { LogScannerService } from './log-insights/log-scanner-service';
 import { createInitialQueryInfo } from './run-queries-shared';
 import { LegacyQueryRunner } from './legacy-query-server/legacyRunner';
+import { NewQueryRunner } from './query-server/query-runner';
 import { QueryRunner } from './queryRunner';
 import { VariantAnalysisView } from './remote-queries/variant-analysis-view';
 import { VariantAnalysisViewSerializer } from './remote-queries/variant-analysis-view-serializer';
@@ -406,6 +408,8 @@ export async function activate(ctx: ExtensionContext): Promise<CodeQLExtensionIn
   return codeQlExtension;
 }
 
+const PACK_GLOBS = ['**/codeql-pack.yml', '**/qlpack.yml', '**/queries.xml', '**/codeql-pack.lock.yml', '**/qlpack.lock.yml', '.codeqlmanifest.json', 'codeql-workspace.yml'];
+
 async function activateWithInstalledDistribution(
   ctx: ExtensionContext,
   distributionManager: DistributionManager,
@@ -436,6 +440,15 @@ async function activateWithInstalledDistribution(
   void logger.log('Initializing query server client.');
   const qs = await createQueryServer(qlConfigurationListener, cliServer, ctx);
 
+
+  for (const glob of PACK_GLOBS) {
+    const fsWatcher = workspace.createFileSystemWatcher(glob);
+    ctx.subscriptions.push(fsWatcher);
+    fsWatcher.onDidChange(async (_uri) => {
+      await qs.clearPackCache();
+    });
+  }
+
   void logger.log('Initializing database manager.');
   const dbm = new DatabaseManager(ctx, qs, cliServer, logger);
   ctx.subscriptions.push(dbm);
@@ -1199,15 +1212,28 @@ async function createQueryServer(qlConfigurationListener: QueryServerConfigListe
     { title: 'CodeQL query server', location: ProgressLocation.Window },
     task
   );
-  const qs = new qsClient.QueryServerClient(
-    qlConfigurationListener,
-    cliServer,
-    qsOpts,
-    progressCallback
-  );
-  ctx.subscriptions.push(qs);
-  await qs.startQueryServer();
-  return new LegacyQueryRunner(qs);
+  if (await cliServer.cliConstraints.supportsNewQueryServer()) {
+    const qs = new newQueryServer.QueryServerClient(
+      qlConfigurationListener,
+      cliServer,
+      qsOpts,
+      progressCallback
+    );
+    ctx.subscriptions.push(qs);
+    await qs.startQueryServer();
+    return new NewQueryRunner(qs);
+
+  } else {
+    const qs = new legacyQueryServer.QueryServerClient(
+      qlConfigurationListener,
+      cliServer,
+      qsOpts,
+      progressCallback
+    );
+    ctx.subscriptions.push(qs);
+    await qs.startQueryServer();
+    return new LegacyQueryRunner(qs);
+  }
 }
 
 function getContextStoragePath(ctx: ExtensionContext) {

extensions/ql-vscode/src/legacy-query-server/legacyRunner.ts

@@ -56,4 +56,10 @@ export class LegacyQueryRunner extends QueryRunner {
   async upgradeDatabaseExplicit(dbItem: DatabaseItem, progress: ProgressCallback, token: CancellationToken): Promise<void> {
     await upgradeDatabaseExplicit(this.qs, dbItem, progress, token);
   }
+
+  async clearPackCache(): Promise<void> {
+    /**
+     * Nothing needs to be done
+     */
+  }
 }

extensions/ql-vscode/src/legacy-query-server/run-queries.ts

@@ -497,11 +497,19 @@ function createSyntheticResult(
   return {
     query: query.queryEvalInfo,
     message,
+    result:{
+      evaluationTime:0,
+      queryId: 0,
+      resultType: messages.QueryResultType.OTHER_ERROR,
+      message,
+      runId: 0,
+    },
     sucessful: false,
     dispose: () => { /**/ },
   };
 }
 
+
 function createSimpleTemplates(templates: Record<string, string> | undefined): messages.TemplateDefinitions | undefined {
   if (!templates) {
     return undefined;

extensions/ql-vscode/src/legacy-query-server/upgrades.ts

@@ -3,7 +3,7 @@ import { getOnDiskWorkspaceFolders, showAndLogErrorMessage, tmpDir } from '../he
 import { ProgressCallback, UserCancellationException } from '../commandRunner';
 import { logger } from '../logging';
 import * as messages from '../pure/legacy-messages';
-import * as qsClient from '../legacy-query-server/queryserver-client';
+import * as qsClient from './queryserver-client';
 import * as tmp from 'tmp-promise';
 import * as path from 'path';
 import { DatabaseItem } from '../databases';

extensions/ql-vscode/src/pure/new-messages.ts

@@ -0,0 +1,215 @@
+/**
+ * Types for messages exchanged during jsonrpc communication with the
+ * the CodeQL query server.
+ *
+ * This file exists in the queryserver and in the vscode extension, and
+ * should be kept in sync between them.
+ *
+ * A note about the namespaces below, which look like they are
+ * essentially enums, namely Severity, ResultColumnKind, and
+ * QueryResultType. By design, for the sake of extensibility, clients
+ * receiving messages of this protocol are supposed to accept any
+ * number for any of these types. We commit to the given meaning of
+ * the numbers listed in constants in the namespaces, and we commit to
+ * the fact that any unknown QueryResultType value counts as an error.
+ */
+
+import * as rpc from 'vscode-jsonrpc';
+import * as shared from './messages-shared';
+
+
+
+/**
+ * Parameters to clear the cache
+ */
+export interface ClearCacheParams {
+  /**
+   * The dataset for which we want to clear the cache
+   */
+  db: string;
+  /**
+   * Whether the cache should actually be cleared.
+   */
+  dryRun: boolean;
+}
+
+/**
+ * Parameters for trimming the cache of a dataset
+ */
+export interface TrimCacheParams {
+  /**
+   * The dataset that we want to trim the cache of.
+   */
+  db: string;
+}
+
+/**
+ * The result of trimming or clearing the cache.
+ */
+export interface ClearCacheResult {
+  /**
+   * A user friendly message saying what was or would be
+   * deleted.
+   */
+  deletionMessage: string;
+}
+
+
+export type QueryResultType = number;
+/**
+ * The result of running a query. This namespace is intentionally not
+ * an enum, see "for the sake of extensibility" comment above.
+ */
+// eslint-disable-next-line @typescript-eslint/no-namespace
+export namespace QueryResultType {
+  /**
+   * The query ran successfully
+   */
+  export const SUCCESS = 0;
+  /**
+   * The query failed due to an reason
+   * that isn't listed
+   */
+  export const OTHER_ERROR = 1;
+  /**
+   * The query failed do to compilation erorrs
+   */
+  export const COMPILATION_ERROR = 2;
+  /**
+    * The query failed due to running out of
+    * memory
+    */
+  export const OOM = 3;
+  /**
+   * The query failed because it was cancelled.
+   */
+  export const CANCELLATION = 4;
+  /**
+   * The dbscheme basename was not the same
+   */
+  export const DBSCHEME_MISMATCH_NAME = 5;
+  /**
+   * No upgrade was found
+   */
+  export const DBSCHEME_NO_UPGRADE = 6;
+}
+
+
+export interface RegisterDatabasesParams {
+  databases: string[];
+}
+
+export interface DeregisterDatabasesParams {
+  databases: string[];
+}
+
+export type RegisterDatabasesResult = {
+  registeredDatabases: string[];
+};
+
+export type DeregisterDatabasesResult = {
+  registeredDatabases: string[];
+};
+
+
+export interface RunQueryParams {
+  /**
+   * The path of the query
+   */
+  queryPath: string,
+  /**
+   * The output path
+   */
+  outputPath: string,
+  /**
+   * The database path
+   */
+  db: string,
+  additionalPacks: string[],
+  target: CompilationTarget,
+  externalInputs: Record<string, string>,
+  singletonExternalInputs: Record<string, string>,
+  dilPath?: string,
+  logPath?: string
+}
+
+export interface RunQueryResult {
+  resultType: QueryResultType,
+  message?: string,
+  expectedDbschemeName?: string,
+  evaluationTime: number;
+}
+
+
+
+export interface UpgradeParams {
+  db: string,
+  additionalPacks: string[],
+}
+
+export type UpgradeResult = Record<string, unknown>;
+
+export type ClearPackCacheParams = Record<string, unknown>;
+export type ClearPackCacheResult = Record<string, unknown>;
+
+/**
+ * A position within a QL file.
+ */
+export type Position = shared.Position;
+
+/**
+ * The way of compiling the query, as a normal query
+ * or a subset of it. Note that precisely one of the two options should be set.
+ */
+export type CompilationTarget = shared.CompilationTarget;
+
+export type QuickEvalOptions = shared.QuickEvalOptions;
+
+export type WithProgressId<T> = shared.WithProgressId<T>;
+export type ProgressMessage = shared.ProgressMessage;
+
+/**
+ * Clear the cache of a dataset
+ */
+export const clearCache = new rpc.RequestType<WithProgressId<ClearCacheParams>, ClearCacheResult, void, void>('evaluation/clearCache');
+/**
+ * Trim the cache of a dataset
+ */
+export const trimCache = new rpc.RequestType<WithProgressId<TrimCacheParams>, ClearCacheResult, void, void>('evaluation/trimCache');
+
+/**
+ * Clear the pack cache
+ */
+export const clearPackCache = new rpc.RequestType<WithProgressId<ClearPackCacheParams>, ClearPackCacheResult, void, void>('evaluation/clearPackCache');
+
+/**
+ * Run a query on a database
+ */
+export const runQuery = new rpc.RequestType<WithProgressId<RunQueryParams>, RunQueryResult, void, void>('evaluation/runQuery');
+
+export const registerDatabases = new rpc.RequestType<
+  WithProgressId<RegisterDatabasesParams>,
+  RegisterDatabasesResult,
+  void,
+  void
+>('evaluation/registerDatabases');
+
+export const deregisterDatabases = new rpc.RequestType<
+  WithProgressId<DeregisterDatabasesParams>,
+  DeregisterDatabasesResult,
+  void,
+  void
+>('evaluation/deregisterDatabases');
+
+
+export const upgradeDatabase = new rpc.RequestType<
+  WithProgressId<UpgradeParams>,
+  UpgradeResult,
+  void,
+  void
+>('evaluation/runUpgrade');
+
+/**
+ * A notification that the progress has been changed.
+ */
+export const progress = shared.progress;

extensions/ql-vscode/src/query-results.ts

@@ -51,7 +51,7 @@ export class CompletedQueryInfo implements QueryWithResults {
   /**
    * The legacy result. This is only set when loading from the query history.
    */
-  readonly result?: legacyMessages.EvaluationResult;
+  readonly result: legacyMessages.EvaluationResult;
   readonly logFileLocation?: string;
   resultCount: number;
 
@@ -83,9 +83,8 @@ export class CompletedQueryInfo implements QueryWithResults {
   ) {
     this.query = evaluation.query;
     this.logFileLocation = evaluation.logFileLocation;
-    if (evaluation.result) {
-      this.result = evaluation.result;
-    }
+    this.result = evaluation.result;
+
     this.message = evaluation.message;
     this.sucessful = evaluation.sucessful;
     // Use the dispose method from the evaluation.

extensions/ql-vscode/src/query-server/query-runner.ts

@@ -0,0 +1,81 @@
+import { CancellationToken } from 'vscode';
+import { ProgressCallback, UserCancellationException } from '../commandRunner';
+import { DatabaseItem } from '../databases';
+import { clearCache, ClearCacheParams, clearPackCache, deregisterDatabases, registerDatabases, upgradeDatabase } from '../pure/new-messages';
+import { InitialQueryInfo, LocalQueryInfo } from '../query-results';
+import { QueryRunner } from '../queryRunner';
+import { QueryWithResults } from '../run-queries-shared';
+import { QueryServerClient } from './queryserver-client';
+import { compileAndRunQueryAgainstDatabase } from './run-queries';
+import * as vscode from 'vscode';
+import { getOnDiskWorkspaceFolders } from '../helpers';
+export class NewQueryRunner extends QueryRunner {
+
+
+  constructor(public readonly qs: QueryServerClient) {
+    super();
+  }
+
+  get cliServer() {
+    return this.qs.cliServer;
+  }
+
+  async restartQueryServer(progress: ProgressCallback, token: CancellationToken): Promise<void> {
+    await this.qs.restartQueryServer(progress, token);
+  }
+
+  onStart(callBack: (progress: ProgressCallback, token: CancellationToken) => Promise<void>) {
+    this.qs.onDidStartQueryServer(callBack);
+  }
+
+  async clearCacheInDatabase(dbItem: DatabaseItem, progress: ProgressCallback, token: CancellationToken): Promise<void> {
+    if (dbItem.contents === undefined) {
+      throw new Error('Can\'t clear the cache in an invalid database.');
+    }
+
+    const db = dbItem.databaseUri.fsPath;
+    const params: ClearCacheParams = {
+      dryRun: false,
+      db,
+    };
+    await this.qs.sendRequest(clearCache, params, token, progress);
+  }
+  async compileAndRunQueryAgainstDatabase(dbItem: DatabaseItem, initialInfo: InitialQueryInfo, queryStorageDir: string, progress: ProgressCallback, token: CancellationToken, templates?: Record<string, string>, queryInfo?: LocalQueryInfo): Promise<QueryWithResults> {
+    return await compileAndRunQueryAgainstDatabase(this.qs.cliServer, this.qs, dbItem, initialInfo, queryStorageDir, progress, token, templates, queryInfo);
+  }
+
+  async deregisterDatabase(progress: ProgressCallback, token: CancellationToken, dbItem: DatabaseItem): Promise<void> {
+    if (dbItem.contents && (await this.qs.cliServer.cliConstraints.supportsDatabaseRegistration())) {
+      const databases: string[] = [dbItem.databaseUri.fsPath];
+      await this.qs.sendRequest(deregisterDatabases, { databases }, token, progress);
+    }
+  }
+  async registerDatabase(progress: ProgressCallback, token: CancellationToken, dbItem: DatabaseItem): Promise<void> {
+    if (dbItem.contents && (await this.qs.cliServer.cliConstraints.supportsDatabaseRegistration())) {
+      const databases: string[] = [dbItem.databaseUri.fsPath];
+      await this.qs.sendRequest(registerDatabases, { databases }, token, progress);
+    }
+  }
+
+  async clearPackCache(): Promise<void> {
+    await this.qs.sendRequest(clearPackCache, {});
+  }
+
+  async upgradeDatabaseExplicit(dbItem: DatabaseItem, progress: ProgressCallback, token: CancellationToken): Promise<void> {
+
+    const yesItem = { title: 'Yes', isCloseAffordance: false };
+    const noItem = { title: 'No', isCloseAffordance: true };
+    const dialogOptions: vscode.MessageItem[] = [yesItem, noItem];
+
+
+
+    const message = `Should the database ${dbItem.databaseUri.fsPath} be destructively upgraded?\n\nThis should not be necessary to run queries 
+    as we will non-destructively update it anyway.`;
+    const chosenItem = await vscode.window.showInformationMessage(message, { modal: true }, ...dialogOptions);
+
+    if (chosenItem !== yesItem) {
+      throw new UserCancellationException('User cancelled the database upgrade.');
+    }
+    await this.qs.sendRequest(upgradeDatabase, { db: dbItem.databaseUri.fsPath, additionalPacks: getOnDiskWorkspaceFolders() }, token, progress);
+  }
+}

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

@@ -0,0 +1,205 @@
+import * as path from 'path';
+import * as fs from 'fs-extra';
+
+import { DisposableObject } from '../pure/disposable-object';
+import { CancellationToken, commands } from 'vscode';
+import { createMessageConnection, RequestType } from 'vscode-jsonrpc';
+import * as cli from '../cli';
+import { QueryServerConfig } from '../config';
+import { Logger, ProgressReporter } from '../logging';
+import { progress, ProgressMessage, WithProgressId } from '../pure/new-messages';
+import * as messages from '../pure/new-messages';
+import { ProgressCallback, ProgressTask } from '../commandRunner';
+import { findQueryLogFile } from '../run-queries-shared';
+import { ServerProcess } from '../json-rpc-server';
+
+type ServerOpts = {
+  logger: Logger;
+  contextStoragePath: string;
+}
+
+
+type WithProgressReporting = (task: (progress: ProgressReporter, token: CancellationToken) => Thenable<void>) => Thenable<void>;
+
+/**
+ * Client that manages a query server process.
+ * The server process is started upon initialization and tracked during its lifetime.
+ * The server process is disposed when the client is disposed, or if the client asks
+ * to restart it (which disposes the existing process and starts a new one).
+ */
+export class QueryServerClient extends DisposableObject {
+
+  serverProcess?: ServerProcess;
+  progressCallbacks: { [key: number]: ((res: ProgressMessage) => void) | undefined };
+  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 activeQueryLogFile: string | undefined;
+
+  constructor(
+    readonly config: QueryServerConfig,
+    readonly cliServer: cli.CodeQLCliServer,
+    readonly opts: ServerOpts,
+    withProgressReporting: WithProgressReporting
+  ) {
+    super();
+    // When the query server configuration changes, restart the query server.
+    if (config.onDidChangeConfiguration !== undefined) {
+      this.push(config.onDidChangeConfiguration(() =>
+        commands.executeCommand('codeQL.restartQueryServer')));
+    }
+    this.withProgressReporting = withProgressReporting;
+    this.nextCallback = 0;
+    this.nextProgress = 0;
+    this.progressCallbacks = {};
+  }
+
+  get logger(): Logger {
+    return this.opts.logger;
+  }
+
+  /** Stops the query server by disposing of the current server process. */
+  private stopQueryServer(): void {
+    if (this.serverProcess !== undefined) {
+      this.disposeAndStopTracking(this.serverProcess);
+    } else {
+      void this.logger.log('No server process to be stopped.');
+    }
+  }
+
+  /** Restarts the query server by disposing of the current server process and then starting a new one. */
+  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 {
+    this.logger.show();
+  }
+
+  /** Starts a new query server process, sending progress messages to the status bar. */
+  async startQueryServer(): Promise<void> {
+    // Use an arrow function to preserve the value of `this`.
+    return this.withProgressReporting((progress, _) => this.startQueryServerImpl(progress));
+  }
+
+  /** Starts a new query server process, sending progress messages to the given reporter. */
+  private async startQueryServerImpl(progressReporter: ProgressReporter): Promise<void> {
+    void this.logger.log('Starting NEW query server.');
+
+    const ramArgs = await this.cliServer.resolveRam(this.config.queryMemoryMb, progressReporter);
+    const args = ['--threads', this.config.numThreads.toString()].concat(ramArgs);
+
+    if (this.config.saveCache) {
+      args.push('--save-cache');
+    }
+
+    if (this.config.cacheSize > 0) {
+      args.push('--max-disk-cache');
+      args.push(this.config.cacheSize.toString());
+    }
+
+    const structuredLogFile = `${this.opts.contextStoragePath}/structured-evaluator-log.json`;
+    await fs.ensureFile(structuredLogFile);
+
+    args.push('--evaluator-log');
+    args.push(structuredLogFile);
+
+    // We hard-code the verbosity level to 5 and minify to false.
+    // This will be the behavior of the per-query structured logging in the CLI after 2.8.3.
+    args.push('--evaluator-log-level');
+    args.push('5');
+
+
+    if (this.config.debug) {
+      args.push('--debug', '--tuple-counting');
+    }
+
+    if (cli.shouldDebugQueryServer()) {
+      args.push('-J=-agentlib:jdwp=transport=dt_socket,address=localhost:9010,server=y,suspend=y,quiet=y');
+    }
+
+    const child = cli.spawnServer(
+      this.config.codeQlPath,
+      'CodeQL query server',
+      ['execute', 'query-server2'],
+      args,
+      this.logger,
+      data => this.logger.log(data.toString(), {
+        trailingNewline: false,
+        additionalLogLocation: this.activeQueryLogFile
+      }),
+      undefined, // no listener for stdout
+      progressReporter
+    );
+    progressReporter.report({ message: 'Connecting to CodeQL query server' });
+    const connection = createMessageConnection(child.stdout, child.stdin);
+    connection.onNotification(progress, res => {
+      const callback = this.progressCallbacks[res.id];
+      if (callback) {
+        callback(res);
+      }
+    });
+    this.serverProcess = new ServerProcess(child, connection, 'Query Server 2', this.logger);
+    // Ensure the server process is disposed together with this client.
+    this.track(this.serverProcess);
+    connection.listen();
+    progressReporter.report({ message: 'Connected to CodeQL query server v2' });
+    this.nextCallback = 0;
+    this.nextProgress = 0;
+    this.progressCallbacks = {};
+  }
+
+  get serverProcessPid(): number {
+    return this.serverProcess!.child.pid || 0;
+  }
+
+  async sendRequest<P, R, E, RO>(type: RequestType<WithProgressId<P>, R, E, RO>, parameter: P, token?: CancellationToken, progress?: (res: ProgressMessage) => void): Promise<R> {
+    const id = this.nextProgress++;
+    this.progressCallbacks[id] = progress;
+
+    this.updateActiveQuery(type.method, parameter);
+    try {
+      if (this.serverProcess === undefined) {
+        throw new Error('No query server process found.');
+      }
+      return await this.serverProcess.connection.sendRequest(type, { body: parameter, progressId: id }, token);
+    } finally {
+      delete this.progressCallbacks[id];
+    }
+  }
+
+  /**
+   * Updates the active query every time there is a new request to compile.
+   * The active query is used to specify the side log.
+   *
+   * This isn't ideal because in situations where there are queries running
+   * in parallel, each query's log messages are interleaved. Fixing this
+   * properly will require a change in the query server.
+   */
+  private updateActiveQuery(method: string, parameter: any): void {
+    if (method === messages.runQuery.method) {
+      this.activeQueryLogFile = findQueryLogFile(path.dirname(path.dirname((parameter as messages.RunQueryParams).outputPath)));
+    }
+  }
+}

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

@@ -0,0 +1,143 @@
+import * as path from 'path';
+import {
+  CancellationToken
+} from 'vscode';
+import * as cli from '../cli';
+import { ProgressCallback } from '../commandRunner';
+import { DatabaseItem } from '../databases';
+import {
+  getOnDiskWorkspaceFolders,
+  showAndLogErrorMessage,
+  showAndLogWarningMessage,
+  tryGetQueryMetadata
+} from '../helpers';
+import { logger } from '../logging';
+import * as messages from '../pure/new-messages';
+import * as legacyMessages from '../pure/legacy-messages';
+import { InitialQueryInfo, LocalQueryInfo } from '../query-results';
+import { QueryEvaluationInfo, QueryWithResults } from '../run-queries-shared';
+import * as qsClient from './queryserver-client';
+
+
+/**
+ * run-queries.ts
+ * --------------
+ *
+ * Compiling and running QL queries.
+ */
+
+
+/**
+ * A collection of evaluation-time information about a query,
+ * including the query itself, and where we have decided to put
+ * temporary files associated with it, such as the compiled query
+ * output and results.
+ */
+
+export async function compileAndRunQueryAgainstDatabase(
+  cliServer: cli.CodeQLCliServer,
+  qs: qsClient.QueryServerClient,
+  dbItem: DatabaseItem,
+  initialInfo: InitialQueryInfo,
+  queryStorageDir: string,
+  progress: ProgressCallback,
+  token: CancellationToken,
+  templates?: Record<string, string>,
+  queryInfo?: LocalQueryInfo, // May be omitted for queries not initiated by the user. If omitted we won't create a structured log for the query.
+): Promise<QueryWithResults> {
+  if (!dbItem.contents || !dbItem.contents.dbSchemeUri) {
+    throw new Error(`Database ${dbItem.databaseUri} does not have a CodeQL database scheme.`);
+  }
+
+  // Read the query metadata if possible, to use in the UI.
+  const metadata = await tryGetQueryMetadata(cliServer, initialInfo.queryPath);
+
+  const hasMetadataFile = (await dbItem.hasMetadataFile());
+  const query = new QueryEvaluationInfo(
+    path.join(queryStorageDir, initialInfo.id),
+    dbItem.databaseUri.fsPath,
+    hasMetadataFile,
+    initialInfo.quickEvalPosition,
+    metadata,
+  );
+
+  if (!dbItem.contents || dbItem.error) {
+    throw new Error('Can\'t run query on invalid database.');
+  }
+  const target = query.quickEvalPosition ? {
+    quickEval: { quickEvalPos: query.quickEvalPosition }
+  } : { query: {} };
+
+  const diskWorkspaceFolders = getOnDiskWorkspaceFolders();
+  const db = dbItem.databaseUri.fsPath;
+  const logPath = queryInfo ? query.evalLogPath : undefined;
+  const queryToRun: messages.RunQueryParams = {
+    db,
+    additionalPacks: diskWorkspaceFolders,
+    externalInputs: {},
+    singletonExternalInputs: templates || {},
+    outputPath: query.resultsPaths.resultsPath,
+    queryPath: initialInfo.queryPath,
+    logPath,
+    target,
+  };
+  await query.createTimestampFile();
+  let result: messages.RunQueryResult | undefined;
+  try {
+    result = await qs.sendRequest(messages.runQuery, queryToRun, token, progress);
+    if (qs.config.customLogDirectory) {
+      void showAndLogWarningMessage(
+        `Custom log directories are no longer supported. The "codeQL.runningQueries.customLogDirectory" setting is deprecated. Unset the setting to stop seeing this message. Query logs saved to ${query.logPath}.`
+      );
+    }
+  } finally {
+    if (queryInfo) {
+      if (await query.hasEvalLog()) {
+        await query.addQueryLogs(queryInfo, qs.cliServer, qs.logger);
+      } else {
+        void showAndLogWarningMessage(`Failed to write structured evaluator log to ${query.evalLogPath}.`);
+      }
+    }
+  }
+
+  if (result.resultType !== messages.QueryResultType.SUCCESS) {
+    const message = result.message || 'Failed to run query';
+    void logger.log(message);
+    void showAndLogErrorMessage(message);
+  }
+  let message;
+  switch (result.resultType) {
+    case messages.QueryResultType.CANCELLATION:
+      message = `cancelled after ${Math.round(result.evaluationTime / 1000)} seconds`;
+      break;
+    case messages.QueryResultType.OOM:
+      message = 'out of memory';
+      break;
+    case messages.QueryResultType.SUCCESS:
+      message = `finished in ${Math.round(result.evaluationTime / 1000)} seconds`;
+      break;
+    case messages.QueryResultType.COMPILATION_ERROR:
+      message = `compilation failed: ${result.message}`;
+      break;
+    case messages.QueryResultType.OTHER_ERROR:
+    default:
+      message = result.message ? `failed: ${result.message}` : 'failed';
+      break;
+  }
+  const sucessful = result.resultType === messages.QueryResultType.SUCCESS;
+  return {
+    query,
+    result: {
+      evaluationTime: result.evaluationTime,
+      queryId: 0,
+      resultType: sucessful ? legacyMessages.QueryResultType.SUCCESS : legacyMessages.QueryResultType.OTHER_ERROR,
+      runId: 0,
+      message
+    },
+    message,
+    sucessful,
+    dispose: () => {
+      qs.logger.removeAdditionalLogLocation(undefined);
+    }
+  };
+}

extensions/ql-vscode/src/queryRunner.ts

@@ -45,4 +45,6 @@ export abstract class QueryRunner {
     progress: ProgressCallback,
     token: CancellationToken,
   ): Promise<void>
+
+  abstract clearPackCache(): Promise<void>
 }

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

@@ -194,9 +194,17 @@ export class QueryEvaluationInfo {
     }
     const compiledQuery = path.join(this.querySaveDir, 'compiledQuery.qlo');
     if (!(await fs.pathExists(compiledQuery))) {
-      throw new Error(
-        `Cannot create DIL because compiled query is missing. ${compiledQuery}`
-      );
+      if (await cliServer.cliConstraints.supportsNewQueryServer()) {
+        // This could be from the new query server
+        // in which case we expect the qlo to be missing so we should ignore it
+        throw new Error(
+          `DIL was not found. ${compiledQuery}`
+        );
+      } else {
+        throw new Error(
+          `Cannot create DIL because compiled query is missing. ${compiledQuery}`
+        );
+      }
     }
 
     await cliServer.generateDil(compiledQuery, this.dilPath);
@@ -210,6 +218,9 @@ export class QueryEvaluationInfo {
     return fs.pathExists(this.evalLogPath);
   }
 
+  /**
+   * Add the structured evaluator log to the query evaluation info.
+   */
   async addQueryLogs(queryInfo: LocalQueryInfo, cliServer: CodeQLCliServer, logger: Logger) {
     queryInfo.evalLogLocation = this.evalLogPath;
     queryInfo.evalLogSummaryLocation = await this.generateHumanReadableLogSummary(cliServer);
@@ -364,7 +375,7 @@ export interface QueryWithResults {
   readonly dispose: () => void;
   readonly sucessful?: boolean;
   readonly message?: string;
-  readonly result?: legacyMessages.EvaluationResult
+  readonly result: legacyMessages.EvaluationResult;
 }
 
 

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

@@ -0,0 +1,201 @@
+import { expect } from 'chai';
+import * as path from 'path';
+import * as tmp from 'tmp';
+import { CancellationTokenSource } from 'vscode-jsonrpc';
+import * as messages from '../../pure/new-messages';
+import * as qsClient from '../../query-server/queryserver-client';
+import * as cli from '../../cli';
+import { CellValue } from '../../pure/bqrs-cli-types';
+import { extensions, Uri } from 'vscode';
+import { CodeQLExtensionInterface } from '../../extension';
+import { fail } from 'assert';
+import { skipIfNoCodeQL } from '../ensureCli';
+import { QueryServerClient } from '../../query-server/queryserver-client';
+import { logger, ProgressReporter } from '../../logging';
+import { QueryResultType } from '../../pure/new-messages';
+import { cleanDatabases, dbLoc, storagePath } from './global.helper';
+import { importArchiveDatabase } from '../../databaseFetcher';
+
+
+const baseDir = path.join(__dirname, '../../../test/data');
+
+const tmpDir = tmp.dirSync({ prefix: 'query_test_', keep: false, unsafeCleanup: true });
+
+const RESULTS_PATH = path.join(tmpDir.name, 'results.bqrs');
+
+const source = new CancellationTokenSource();
+const token = source.token;
+
+class Checkpoint<T> {
+  private res: () => void;
+  private rej: (e: Error) => void;
+  private promise: Promise<T>;
+
+  constructor() {
+    this.res = () => { /**/ };
+    this.rej = () => { /**/ };
+    this.promise = new Promise((res, rej) => {
+      this.res = res as () => Record<string, never>;
+      this.rej = rej;
+    });
+  }
+
+  async done(): Promise<T> {
+    return this.promise;
+  }
+
+  async resolve(): Promise<void> {
+    await (this.res)();
+  }
+
+  async reject(e: Error): Promise<void> {
+    await (this.rej)(e);
+  }
+}
+
+type ResultSets = {
+  [name: string]: CellValue[][];
+}
+
+type QueryTestCase = {
+  queryPath: string;
+  expectedResultSets: ResultSets;
+}
+
+// Test cases: queries to run and their expected results.
+const queryTestCases: QueryTestCase[] = [
+  {
+    queryPath: path.join(baseDir, 'query.ql'),
+    expectedResultSets: {
+      '#select': [[42, 3.14159, 'hello world', true]]
+    }
+  },
+  {
+    queryPath: path.join(baseDir, 'compute-default-strings.ql'),
+    expectedResultSets: {
+      '#select': [[{ label: '(no string representation)' }]]
+    }
+  },
+  {
+    queryPath: path.join(baseDir, 'multiple-result-sets.ql'),
+    expectedResultSets: {
+      'edges': [[1, 2], [2, 3]],
+      '#select': [['s']]
+    }
+  }
+];
+
+const nullProgressReporter: ProgressReporter = { report: () => { /** ignore */ } };
+
+describe('using the new query server', function() {
+  before(function() {
+    skipIfNoCodeQL(this);
+  });
+
+  // Note this does not work with arrow functions as the test case bodies:
+  // ensure they are all written with standard anonymous functions.
+  this.timeout(20000);
+
+  let qs: qsClient.QueryServerClient;
+  let cliServer: cli.CodeQLCliServer;
+  let db: string;
+  before(async () => {
+    try {
+      const extension = await extensions.getExtension<CodeQLExtensionInterface | Record<string, never>>('GitHub.vscode-codeql')!.activate();
+      if ('cliServer' in extension && 'databaseManager' in extension) {
+        cliServer = extension.cliServer;
+
+        cliServer.quiet = true;
+        if (!(await cliServer.cliConstraints.supportsNewQueryServer())) {
+          this.ctx.skip();
+        }
+        qs = new QueryServerClient({
+          codeQlPath: (await extension.distributionManager.getCodeQlPathWithoutVersionCheck()) || '',
+          debug: false,
+          cacheSize: 0,
+          numThreads: 1,
+          saveCache: false,
+          timeoutSecs: 0
+        }, cliServer, {
+          contextStoragePath: tmpDir.name,
+          logger
+        }, task => task(nullProgressReporter, new CancellationTokenSource().token));
+        await qs.startQueryServer();
+
+        // Unlike the old query sevre the new one wants a database and the empty direcrtory is not valid.
+        // Add a database, but make sure the database manager is empty first
+        await cleanDatabases(extension.databaseManager);
+        const uri = Uri.file(dbLoc);
+        const maybeDbItem = await importArchiveDatabase(
+          uri.toString(true),
+          extension.databaseManager,
+          storagePath,
+          () => { /**ignore progress */ },
+          token,
+        );
+
+        if (!maybeDbItem) {
+          throw new Error('Could not import database');
+        }
+        db = maybeDbItem.databaseUri.fsPath;
+      } else {
+        throw new Error('Extension not initialized. Make sure cli is downloaded and installed properly.');
+      }
+    } catch (e) {
+      fail(e as Error);
+    }
+
+  });
+
+  for (const queryTestCase of queryTestCases) {
+    const queryName = path.basename(queryTestCase.queryPath);
+    const evaluationSucceeded = new Checkpoint<void>();
+    const parsedResults = new Checkpoint<void>();
+
+    it('should register the database', async () => {
+      await qs.sendRequest(messages.registerDatabases, { databases: [db] }, token, (() => { /**/ }) as any);
+    });
+
+
+    it(`should be able to run query ${queryName}`, async function() {
+      try {
+        const params: messages.RunQueryParams = {
+          db,
+          queryPath: queryTestCase.queryPath,
+          outputPath: RESULTS_PATH,
+          additionalPacks: [],
+          externalInputs: {},
+          singletonExternalInputs: {},
+          target: { query: {} }
+        };
+        const result = await qs.sendRequest(messages.runQuery, params, token, () => { /**/ });
+        expect(result.resultType).to.equal(QueryResultType.SUCCESS);
+        await evaluationSucceeded.resolve();
+      }
+      catch (e) {
+        await evaluationSucceeded.reject(e as Error);
+      }
+    });
+
+    const actualResultSets: ResultSets = {};
+    it(`should be able to parse results of query ${queryName}`, async function() {
+      await evaluationSucceeded.done();
+      const info = await cliServer.bqrsInfo(RESULTS_PATH);
+
+      for (const resultSet of info['result-sets']) {
+        const decoded = await cliServer.bqrsDecode(RESULTS_PATH, resultSet.name);
+        actualResultSets[resultSet.name] = decoded.tuples;
+      }
+      await parsedResults.resolve();
+    });
+
+    it(`should have correct results for query ${queryName}`, async function() {
+      await parsedResults.done();
+      expect(actualResultSets!).not.to.be.empty;
+      expect(Object.keys(actualResultSets!).sort()).to.eql(Object.keys(queryTestCase.expectedResultSets).sort());
+      for (const name in queryTestCase.expectedResultSets) {
+        expect(actualResultSets![name]).to.eql(queryTestCase.expectedResultSets[name], `Results for query predicate ${name} do not match`);
+      }
+    });
+  }
+});

extensions/ql-vscode/src/vscode-tests/no-workspace/query-history.test.ts

@@ -20,6 +20,7 @@ import { RemoteQueriesManager } from '../../remote-queries/remote-queries-manage
 import { ResultsView } from '../../interface';
 import { EvalLogViewer } from '../../eval-log-viewer';
 import { QueryRunner } from '../../queryRunner';
+import { QueryResultType } from '../../pure/legacy-messages';
 
 describe('query-history', () => {
   const mockExtensionLocation = path.join(tmpDir.name, 'mock-extension-location');
@@ -786,7 +787,13 @@ describe('query-history', () => {
       } as unknown as QueryEvaluationInfo,
       sucessful: didRunSuccessfully,
       message: 'foo',
-      dispose: sandbox.spy()
+      dispose: sandbox.spy(),
+      result: {
+        evaluationTime: 1,
+        queryId: 0,
+        runId: 0,
+        resultType: QueryResultType.SUCCESS,
+      }
     };
   }
 

extensions/ql-vscode/src/vscode-tests/no-workspace/query-results.test.ts

@@ -328,6 +328,12 @@ describe('query-results', () => {
       sucessful: didRunSuccessfully,
       message: 'foo',
       dispose: disposeSpy,
+      result: {
+        evaluationTime: 1,
+        queryId: 0,
+        runId: 0,
+        resultType: QueryResultType.SUCCESS,
+      }
     };
 
     if (includeSpies) {