Merge pull request #3171 from github/koesie10/remove-new-query-runner-abstraction

Remove `QueryRunner` abstraction

Author: koesie10

extensions/ql-vscode/src/extension.ts

@@ -129,7 +129,7 @@ import { ModelEditorModule } from "./model-editor/model-editor-module";
 import { TestManager } from "./query-testing/test-manager";
 import { TestRunner } from "./query-testing/test-runner";
 import { TestManagerBase } from "./query-testing/test-manager-base";
-import { NewQueryRunner, QueryRunner, QueryServerClient } from "./query-server";
+import { QueryRunner, QueryServerClient } from "./query-server";
 import { QueriesModule } from "./queries-panel/queries-module";
 import { OpenReferencedFileCodeLensProvider } from "./local-queries/open-referenced-file-code-lens-provider";
 import { LanguageContextStore } from "./language-context-store";
@@ -1256,7 +1256,7 @@ async function createQueryServer(
   );
   ctx.subscriptions.push(qs);
   await qs.startQueryServer();
-  return new NewQueryRunner(qs);
+  return new QueryRunner(qs);
 }
 
 function getContextStoragePath(ctx: ExtensionContext) {

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

@@ -1,4 +1,3 @@
-export * from "./new-query-runner";
 export * from "./query-runner";
 export * from "./query-server-client";
 export * from "./run-queries";

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

@@ -1,12 +1,29 @@
-import { CancellationToken } from "vscode";
+import vscode, { CancellationToken } from "vscode";
 import { CodeQLCliServer } from "../codeql-cli/cli";
-import { ProgressCallback } from "../common/vscode/progress";
+import {
+  ProgressCallback,
+  UserCancellationException,
+} from "../common/vscode/progress";
 import { DatabaseItem } from "../databases/local-databases";
 import { QueryOutputDir } from "../run-queries-shared";
-import { Position, QueryResultType } from "./new-messages";
+import {
+  clearCache,
+  ClearCacheParams,
+  clearPackCache,
+  deregisterDatabases,
+  Position,
+  QueryResultType,
+  registerDatabases,
+  trimCache,
+  TrimCacheParams,
+  upgradeDatabase,
+} from "./new-messages";
 import { BaseLogger, Logger } from "../common/logging";
 import { basename, join } from "path";
 import { nanoid } from "nanoid";
+import { QueryServerClient } from "./query-server-client";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { compileAndRunQueryAgainstDatabaseCore } from "./run-queries";
 
 export interface CoreQueryTarget {
   /** The full path to the query. */
@@ -45,37 +62,69 @@ export interface CoreQueryRun {
 export type CoreCompletedQuery = CoreQueryResults &
   Omit<CoreQueryRun, "evaluate">;
 
-export abstract class QueryRunner {
-  abstract restartQueryServer(
+export class QueryRunner {
+  constructor(public readonly qs: QueryServerClient) {}
+
+  get cliServer(): CodeQLCliServer {
+    return this.qs.cliServer;
+  }
+
+  get customLogDirectory(): string | undefined {
+    return this.qs.config.customLogDirectory;
+  }
+
+  get logger(): Logger {
+    return this.qs.logger;
+  }
+
+  async restartQueryServer(
     progress: ProgressCallback,
     token: CancellationToken,
-  ): Promise<void>;
-
-  abstract cliServer: CodeQLCliServer;
-  abstract customLogDirectory: string | undefined;
-  abstract logger: Logger;
+  ): Promise<void> {
+    await this.qs.restartQueryServer(progress, token);
+  }
 
-  abstract onStart(
-    arg0: (
+  onStart(
+    callBack: (
       progress: ProgressCallback,
       token: CancellationToken,
     ) => Promise<void>,
-  ): void;
+  ) {
+    this.qs.onDidStartQueryServer(callBack);
+  }
 
-  abstract clearCacheInDatabase(
+  async clearCacheInDatabase(
     dbItem: DatabaseItem,
     token: CancellationToken,
-  ): Promise<void>;
+  ): Promise<void> {
+    if (dbItem.contents === undefined) {
+      throw new Error("Can't clear the cache in an invalid database.");
+    }
 
-  abstract trimCacheInDatabase(
+    const db = dbItem.databaseUri.fsPath;
+    const params: ClearCacheParams = {
+      dryRun: false,
+      db,
+    };
+    await this.qs.sendRequest(clearCache, params, token);
+  }
+
+  async trimCacheInDatabase(
     dbItem: DatabaseItem,
     token: CancellationToken,
-  ): Promise<void>;
+  ): Promise<void> {
+    if (dbItem.contents === undefined) {
+      throw new Error("Can't trim the cache in an invalid database.");
+    }
 
-  /**
-   * Overridden in subclasses to evaluate the query via the query server and return the results.
-   */
-  public abstract compileAndRunQueryAgainstDatabaseCore(
+    const db = dbItem.databaseUri.fsPath;
+    const params: TrimCacheParams = {
+      db,
+    };
+    await this.qs.sendRequest(trimCache, params, token);
+  }
+
+  public async compileAndRunQueryAgainstDatabaseCore(
     dbPath: string,
     query: CoreQueryTarget,
     additionalPacks: string[],
@@ -87,19 +136,72 @@ export abstract class QueryRunner {
     token: CancellationToken,
     templates: Record<string, string> | undefined,
     logger: BaseLogger,
-  ): Promise<CoreQueryResults>;
+  ): Promise<CoreQueryResults> {
+    return await compileAndRunQueryAgainstDatabaseCore(
+      this.qs,
+      dbPath,
+      query,
+      generateEvalLog,
+      additionalPacks,
+      extensionPacks,
+      additionalRunQueryArgs,
+      outputDir,
+      progress,
+      token,
+      templates,
+      logger,
+    );
+  }
 
-  abstract deregisterDatabase(dbItem: DatabaseItem): Promise<void>;
+  async deregisterDatabase(dbItem: DatabaseItem): Promise<void> {
+    if (dbItem.contents) {
+      const databases: string[] = [dbItem.databaseUri.fsPath];
+      await this.qs.sendRequest(deregisterDatabases, { databases });
+    }
+  }
+  async registerDatabase(dbItem: DatabaseItem): Promise<void> {
+    if (dbItem.contents) {
+      const databases: string[] = [dbItem.databaseUri.fsPath];
+      await this.qs.sendRequest(registerDatabases, { databases });
+    }
+  }
 
-  abstract registerDatabase(dbItem: DatabaseItem): Promise<void>;
+  async clearPackCache(): Promise<void> {
+    await this.qs.sendRequest(clearPackCache, {});
+  }
 
-  abstract upgradeDatabaseExplicit(
+  async upgradeDatabaseExplicit(
     dbItem: DatabaseItem,
     progress: ProgressCallback,
     token: CancellationToken,
-  ): Promise<void>;
+  ): Promise<void> {
+    const yesItem = { title: "Yes", isCloseAffordance: false };
+    const noItem = { title: "No", isCloseAffordance: true };
+    const dialogOptions: vscode.MessageItem[] = [yesItem, noItem];
 
-  abstract clearPackCache(): Promise<void>;
+    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,
+    );
+  }
 
   /**
    * Create a `CoreQueryRun` object. This creates an object whose `evaluate()` function can be