Merge pull request #1498 from alexet/alexet/prepare-new-qs

QueryServer: Abstract over the query running parts.

Author: aeisenberg

extensions/ql-vscode/src/cli.ts

@@ -16,7 +16,7 @@ import { DistributionProvider, FindDistributionResultKind } from './distribution
 import { assertNever, getErrorMessage, getErrorStack } from './pure/helpers-pure';
 import { QueryMetadata, SortDirection } from './pure/interface-types';
 import { Logger, ProgressReporter } from './logging';
-import { CompilationMessage } from './pure/messages';
+import { CompilationMessage } from './pure/legacy-messages';
 import { sarifParser } from './sarif-parser';
 import { dbSchemeToLanguage, walkDirectory } from './helpers';
 
@@ -1248,6 +1248,9 @@ export class CliVersionConstraint {
    */
   public static CLI_VERSION_WITH_LANGUAGE = new SemVer('2.4.1');
 
+
+  public static CLI_VERSION_WITH_NONDESTURCTIVE_UPGRADES = new SemVer('2.4.2');
+
   /**
    * CLI version where `codeql resolve upgrades` supports
    * the `--allow-downgrades` flag
@@ -1264,14 +1267,6 @@ export class CliVersionConstraint {
    */
   public static CLI_VERSION_WITH_DB_REGISTRATION = new SemVer('2.4.1');
 
-  /**
-   * CLI version where non destructive upgrades were introduced.
-   *
-   * This was landed in multiple parts so this is the version where all necessary feature were supported.
-   */
-  public static CLI_VERSION_WITH_NON_DESTRUCTIVE_UPGRADES = new SemVer('2.4.2');
-
-
   /**
    * CLI version where the `--allow-library-packs` option to `codeql resolve queries` was
    * introduced.
@@ -1351,6 +1346,10 @@ export class CliVersionConstraint {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_LANGUAGE);
   }
 
+  public async supportsNonDestructiveUpgrades() {
+    return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_NONDESTURCTIVE_UPGRADES);
+  }
+
   public async supportsDowngrades() {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_DOWNGRADES);
   }
@@ -1367,10 +1366,6 @@ export class CliVersionConstraint {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_DB_REGISTRATION);
   }
 
-  async supportsNonDestructiveUpgrades(): Promise<boolean> {
-    return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_NON_DESTRUCTIVE_UPGRADES);
-  }
-
   async supportsDatabaseUnbundle() {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_DATABASE_UNBUNDLE);
   }

extensions/ql-vscode/src/contextual/astBuilder.ts

@@ -1,10 +1,10 @@
-import { QueryWithResults } from '../run-queries';
 import { CodeQLCliServer } from '../cli';
 import { DecodedBqrsChunk, BqrsId, EntityValue } from '../pure/bqrs-cli-types';
 import { DatabaseItem } from '../databases';
 import { ChildAstItem, AstItem } from '../astViewer';
 import fileRangeFromURI from './fileRangeFromURI';
 import { Uri } from 'vscode';
+import { QueryWithResults } from '../run-queries-shared';
 
 /**
  * A class that wraps a tree of QL results from a query that

extensions/ql-vscode/src/contextual/locationFinder.ts

@@ -3,13 +3,12 @@ import { ColumnKindCode, EntityValue, getResultSetSchema, ResultSetSchema } from
 import { CodeQLCliServer } from '../cli';
 import { DatabaseManager, DatabaseItem } from '../databases';
 import fileRangeFromURI from './fileRangeFromURI';
-import * as messages from '../pure/messages';
-import { QueryServerClient } from '../queryserver-client';
-import { QueryWithResults, compileAndRunQueryAgainstDatabase, createInitialQueryInfo } from '../run-queries';
 import { ProgressCallback } from '../commandRunner';
 import { KeyType } from './keyType';
 import { qlpackOfDatabase, resolveQueries } from './queryResolver';
 import { CancellationToken, LocationLink, Uri } from 'vscode';
+import { createInitialQueryInfo, QueryWithResults } from '../run-queries-shared';
+import { QueryRunner } from '../queryRunner';
 
 export const SELECT_QUERY_NAME = '#select';
 export const TEMPLATE_NAME = 'selectedSourceFile';
@@ -35,7 +34,7 @@ export interface FullLocationLink extends LocationLink {
  */
 export async function getLocationsForUriString(
   cli: CodeQLCliServer,
-  qs: QueryServerClient,
+  qs: QueryRunner,
   dbm: DatabaseManager,
   uriString: string,
   keyType: KeyType,
@@ -65,19 +64,8 @@ export async function getLocationsForUriString(
       },
       false
     );
-
-    const results = await compileAndRunQueryAgainstDatabase(
-      cli,
-      qs,
-      db,
-      initialInfo,
-      queryStorageDir,
-      progress,
-      token,
-      templates
-    );
-
-    if (results.result.resultType == messages.QueryResultType.SUCCESS) {
+    const results = await qs.compileAndRunQueryAgainstDatabase(db, initialInfo, queryStorageDir, progress, token, templates);
+    if (results.sucessful) {
       links.push(...await getLinksFromResults(results, cli, db, filter));
     }
   }
@@ -114,15 +102,9 @@ async function getLinksFromResults(
   return localLinks;
 }
 
-function createTemplates(path: string): messages.TemplateDefinitions {
+function createTemplates(path: string): Record<string, string> {
   return {
-    [TEMPLATE_NAME]: {
-      values: {
-        tuples: [[{
-          stringValue: path
-        }]]
-      }
-    }
+    [TEMPLATE_NAME]: path
   };
 }
 

extensions/ql-vscode/src/contextual/templateProvider.ts

@@ -16,16 +16,15 @@ import { CodeQLCliServer } from '../cli';
 import { DatabaseManager } from '../databases';
 import { CachedOperation } from '../helpers';
 import { ProgressCallback, withProgress } from '../commandRunner';
-import * as messages from '../pure/messages';
-import { QueryServerClient } from '../queryserver-client';
-import { compileAndRunQueryAgainstDatabase, createInitialQueryInfo, QueryWithResults } from '../run-queries';
 import AstBuilder from './astBuilder';
 import {
   KeyType,
 } from './keyType';
 import { FullLocationLink, getLocationsForUriString, TEMPLATE_NAME } from './locationFinder';
 import { qlpackOfDatabase, resolveQueries } from './queryResolver';
 import { isCanary, NO_CACHE_AST_VIEWER } from '../config';
+import { createInitialQueryInfo, QueryWithResults } from '../run-queries-shared';
+import { QueryRunner } from '../queryRunner';
 
 /**
  * Run templated CodeQL queries to find definitions and references in
@@ -39,7 +38,7 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
 
   constructor(
     private cli: CodeQLCliServer,
-    private qs: QueryServerClient,
+    private qs: QueryRunner,
     private dbm: DatabaseManager,
     private queryStorageDir: string,
   ) {
@@ -83,7 +82,7 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
 
   constructor(
     private cli: CodeQLCliServer,
-    private qs: QueryServerClient,
+    private qs: QueryRunner,
     private dbm: DatabaseManager,
     private queryStorageDir: string,
   ) {
@@ -137,7 +136,7 @@ export class TemplatePrintAstProvider {
 
   constructor(
     private cli: CodeQLCliServer,
-    private qs: QueryServerClient,
+    private qs: QueryRunner,
     private dbm: DatabaseManager,
     private queryStorageDir: string,
   ) {
@@ -195,14 +194,9 @@ export class TemplatePrintAstProvider {
     }
 
     const query = queries[0];
-    const templates: messages.TemplateDefinitions = {
-      [TEMPLATE_NAME]: {
-        values: {
-          tuples: [[{
-            stringValue: zippedArchive.pathWithinSourceArchive
-          }]]
-        }
-      }
+    const templates: Record<string, string> = {
+      [TEMPLATE_NAME]:
+        zippedArchive.pathWithinSourceArchive
     };
 
     const initialInfo = await createInitialQueryInfo(
@@ -215,9 +209,7 @@ export class TemplatePrintAstProvider {
     );
 
     return {
-      query: await compileAndRunQueryAgainstDatabase(
-        this.cli,
-        this.qs,
+      query: await this.qs.compileAndRunQueryAgainstDatabase(
         db,
         initialInfo,
         this.queryStorageDir,
@@ -231,23 +223,23 @@ export class TemplatePrintAstProvider {
 }
 
 export class TemplatePrintCfgProvider {
-  private cache: CachedOperation<[Uri, messages.TemplateDefinitions] | undefined>;
+  private cache: CachedOperation<[Uri, Record<string, string>] | undefined>;
 
   constructor(
     private cli: CodeQLCliServer,
     private dbm: DatabaseManager,
   ) {
-    this.cache = new CachedOperation<[Uri, messages.TemplateDefinitions] | undefined>(this.getCfgUri.bind(this));
+    this.cache = new CachedOperation<[Uri, Record<string, string>] | undefined>(this.getCfgUri.bind(this));
   }
 
-  async provideCfgUri(document?: TextDocument): Promise<[Uri, messages.TemplateDefinitions] | undefined> {
+  async provideCfgUri(document?: TextDocument): Promise<[Uri, Record<string, string>] | undefined> {
     if (!document) {
       return;
     }
     return await this.cache.get(document.uri.toString());
   }
 
-  private async getCfgUri(uriString: string): Promise<[Uri, messages.TemplateDefinitions]> {
+  private async getCfgUri(uriString: string): Promise<[Uri, Record<string, string>]> {
     const uri = Uri.parse(uriString, true);
     if (uri.scheme !== zipArchiveScheme) {
       throw new Error('CFG Viewing is only available for databases with zipped source archives.');
@@ -275,14 +267,8 @@ export class TemplatePrintCfgProvider {
 
     const queryUri = Uri.file(queries[0]);
 
-    const templates: messages.TemplateDefinitions = {
-      [TEMPLATE_NAME]: {
-        values: {
-          tuples: [[{
-            stringValue: zippedArchive.pathWithinSourceArchive
-          }]]
-        }
-      }
+    const templates: Record<string, string> = {
+      [TEMPLATE_NAME]: zippedArchive.pathWithinSourceArchive
     };
 
     return [queryUri, templates];

extensions/ql-vscode/src/databases-ui.ts

@@ -28,9 +28,6 @@ import {
   showAndLogErrorMessage
 } from './helpers';
 import { logger } from './logging';
-import { clearCacheInDatabase } from './run-queries';
-import * as qsClient from './queryserver-client';
-import { upgradeDatabaseExplicit } from './upgrades';
 import {
   importArchiveDatabase,
   promptImportGithubDatabase,
@@ -40,6 +37,7 @@ import {
 import { CancellationToken } from 'vscode';
 import { asyncFilter, getErrorMessage } from './pure/helpers-pure';
 import { Credentials } from './authentication';
+import { QueryRunner } from './queryRunner';
 import { isCanary } from './config';
 
 type ThemableIconPath = { light: string; dark: string } | string;
@@ -220,7 +218,7 @@ export class DatabaseUI extends DisposableObject {
 
   public constructor(
     private databaseManager: DatabaseManager,
-    private readonly queryServer: qsClient.QueryServerClient | undefined,
+    private readonly queryServer: QueryRunner | undefined,
     private readonly storagePath: string,
     readonly extensionPath: string,
     private readonly getCredentials: () => Promise<Credentials>
@@ -390,12 +388,11 @@ export class DatabaseUI extends DisposableObject {
   handleChooseDatabaseFolder = async (
     progress: ProgressCallback,
     token: CancellationToken
-  ): Promise<DatabaseItem | undefined> => {
+  ): Promise<void> => {
     try {
-      return await this.chooseAndSetDatabase(true, progress, token);
+      await this.chooseAndSetDatabase(true, progress, token);
     } catch (e) {
       void showAndLogErrorMessage(getErrorMessage(e));
-      return undefined;
     }
   };
 
@@ -458,12 +455,11 @@ export class DatabaseUI extends DisposableObject {
   handleChooseDatabaseArchive = async (
     progress: ProgressCallback,
     token: CancellationToken
-  ): Promise<DatabaseItem | undefined> => {
+  ): Promise<void> => {
     try {
-      return await this.chooseAndSetDatabase(false, progress, token);
+      await this.chooseAndSetDatabase(false, progress, token);
     } catch (e) {
       void showAndLogErrorMessage(getErrorMessage(e));
-      return undefined;
     }
   };
 
@@ -576,8 +572,7 @@ export class DatabaseUI extends DisposableObject {
 
     // Search for upgrade scripts in any workspace folders available
 
-    await upgradeDatabaseExplicit(
-      this.queryServer,
+    await this.queryServer.upgradeDatabaseExplicit(
       databaseItem,
       progress,
       token
@@ -592,8 +587,7 @@ export class DatabaseUI extends DisposableObject {
       this.queryServer !== undefined &&
       this.databaseManager.currentDatabaseItem !== undefined
     ) {
-      await clearCacheInDatabase(
-        this.queryServer,
+      await this.queryServer.clearCacheInDatabase(
         this.databaseManager.currentDatabaseItem,
         progress,
         token

extensions/ql-vscode/src/databases.ts

@@ -17,9 +17,8 @@ import {
 import { zipArchiveScheme, encodeArchiveBasePath, decodeSourceArchiveUri, encodeSourceArchiveUri } from './archive-filesystem-provider';
 import { DisposableObject } from './pure/disposable-object';
 import { Logger, logger } from './logging';
-import { registerDatabases, Dataset, deregisterDatabases } from './pure/messages';
-import { QueryServerClient } from './queryserver-client';
 import { getErrorMessage } from './pure/helpers-pure';
+import { QueryRunner } from './queryRunner';
 
 /**
  * databases.ts
@@ -555,13 +554,13 @@ export class DatabaseManager extends DisposableObject {
 
   constructor(
     private readonly ctx: ExtensionContext,
-    private readonly qs: QueryServerClient,
+    private readonly qs: QueryRunner,
     private readonly cli: cli.CodeQLCliServer,
     public logger: Logger
   ) {
     super();
 
-    qs.onDidStartQueryServer(this.reregisterDatabases.bind(this));
+    qs.onStart(this.reregisterDatabases.bind(this));
 
     // Let this run async.
     void this.loadPersistedState();
@@ -860,27 +859,14 @@ export class DatabaseManager extends DisposableObject {
     token: vscode.CancellationToken,
     dbItem: DatabaseItem,
   ) {
-    if (dbItem.contents && (await this.cli.cliConstraints.supportsDatabaseRegistration())) {
-      const databases: Dataset[] = [{
-        dbDir: dbItem.contents.datasetUri.fsPath,
-        workingSet: 'default'
-      }];
-      await this.qs.sendRequest(deregisterDatabases, { databases }, token, progress);
-    }
+    await this.qs.deregisterDatabase(progress, token, dbItem);
   }
-
   private async registerDatabase(
     progress: ProgressCallback,
     token: vscode.CancellationToken,
     dbItem: DatabaseItem,
   ) {
-    if (dbItem.contents && (await this.cli.cliConstraints.supportsDatabaseRegistration())) {
-      const databases: Dataset[] = [{
-        dbDir: dbItem.contents.datasetUri.fsPath,
-        workingSet: 'default'
-      }];
-      await this.qs.sendRequest(registerDatabases, { databases }, token, progress);
-    }
+    await this.qs.registerDatabase(progress, token, dbItem);
   }
 
   private updatePersistedCurrentDatabaseItem(): void {