Add support for quick eval count to the query runner (#2417)

* Add support for quick eval count to the query runner

This only adds support internally to the query runner,
without any UI support.

* Fix some tests

extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts

@@ -78,7 +78,11 @@ export async function runQuery({
 
   const queryRun = queryRunner.createQueryRun(
     databaseItem.databaseUri.fsPath,
-    { queryPath: queryFile, quickEvalPosition: undefined },
+    {
+      queryPath: queryFile,
+      quickEvalPosition: undefined,
+      quickEvalCountOnly: false,
+    },
     false,
     getOnDiskWorkspaceFolders(),
     extensionPacks,

extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts

@@ -53,7 +53,11 @@ async function getModeledMethodsFromFlow(
 
   const queryRun = queryRunner.createQueryRun(
     databaseItem.databaseUri.fsPath,
-    { queryPath: query, quickEvalPosition: undefined },
+    {
+      queryPath: query,
+      quickEvalPosition: undefined,
+      quickEvalCountOnly: false,
+    },
     false,
     getOnDiskWorkspaceFolders(),
     undefined,

extensions/ql-vscode/src/debugger/debug-configuration.ts

@@ -105,7 +105,7 @@ export class QLDebugConfigurationProvider
       validateQueryPath(qlConfiguration.query, quickEval);
 
       const quickEvalContext = quickEval
-        ? await getQuickEvalContext(undefined)
+        ? await getQuickEvalContext(undefined, false)
         : undefined;
 
       const resultConfiguration: QLResolvedDebugConfiguration = {

extensions/ql-vscode/src/debugger/debug-session.ts

@@ -155,6 +155,7 @@ class RunningQuery extends DisposableObject {
       {
         queryPath: config.query,
         quickEvalPosition: quickEvalContext?.quickEvalPosition,
+        quickEvalCountOnly: quickEvalContext?.quickEvalCount,
       },
       true,
       config.additionalPacks,

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

@@ -74,7 +74,7 @@ class QLDebugAdapterTracker
 
   public async quickEval(): Promise<void> {
     const args: CodeQLProtocol.QuickEvalRequest["arguments"] = {
-      quickEvalContext: await getQuickEvalContext(undefined),
+      quickEvalContext: await getQuickEvalContext(undefined, false),
     };
     await this.session.customRequest("codeql-quickeval", args);
   }

extensions/ql-vscode/src/language-support/ast-viewer/ast-cfg-commands.ts

@@ -2,7 +2,7 @@ import { Uri, window } from "vscode";
 import { withProgress } from "../../common/vscode/progress";
 import { AstViewer } from "./ast-viewer";
 import { AstCfgCommands } from "../../common/commands";
-import { LocalQueries } from "../../local-queries";
+import { LocalQueries, QuickEvalType } from "../../local-queries";
 import {
   TemplatePrintAstProvider,
   TemplatePrintCfgProvider,
@@ -47,7 +47,7 @@ export function getAstCfgCommands({
         );
         if (res) {
           await localQueries.compileAndRunQuery(
-            false,
+            QuickEvalType.None,
             res[0],
             progress,
             token,

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

@@ -72,6 +72,12 @@ async function promptToSaveQueryIfNeeded(query: SelectedQuery): Promise<void> {
   }
 }
 
+export enum QuickEvalType {
+  None,
+  QuickEval,
+  QuickEvalCount,
+}
+
 export class LocalQueries extends DisposableObject {
   public constructor(
     private readonly app: App,
@@ -115,7 +121,13 @@ export class LocalQueries extends DisposableObject {
   private async runQuery(uri: Uri | undefined): Promise<void> {
     await withProgress(
       async (progress, token) => {
-        await this.compileAndRunQuery(false, uri, progress, token, undefined);
+        await this.compileAndRunQuery(
+          QuickEvalType.None,
+          uri,
+          progress,
+          token,
+          undefined,
+        );
       },
       {
         title: "Running query",
@@ -185,7 +197,7 @@ export class LocalQueries extends DisposableObject {
         await Promise.all(
           queryUris.map(async (uri) =>
             this.compileAndRunQuery(
-              false,
+              QuickEvalType.None,
               uri,
               wrappedProgress,
               token,
@@ -204,7 +216,13 @@ export class LocalQueries extends DisposableObject {
   private async quickEval(uri: Uri): Promise<void> {
     await withProgress(
       async (progress, token) => {
-        await this.compileAndRunQuery(true, uri, progress, token, undefined);
+        await this.compileAndRunQuery(
+          QuickEvalType.QuickEval,
+          uri,
+          progress,
+          token,
+          undefined,
+        );
       },
       {
         title: "Running query",
@@ -217,7 +235,7 @@ export class LocalQueries extends DisposableObject {
     await withProgress(
       async (progress, token) =>
         await this.compileAndRunQuery(
-          true,
+          QuickEvalType.QuickEval,
           uri,
           progress,
           token,
@@ -331,7 +349,7 @@ export class LocalQueries extends DisposableObject {
   }
 
   public async compileAndRunQuery(
-    quickEval: boolean,
+    quickEval: QuickEvalType,
     queryUri: Uri | undefined,
     progress: ProgressCallback,
     token: CancellationToken,
@@ -352,7 +370,7 @@ export class LocalQueries extends DisposableObject {
 
   /** Used by tests */
   public async compileAndRunQueryInternal(
-    quickEval: boolean,
+    quickEval: QuickEvalType,
     queryUri: Uri | undefined,
     progress: ProgressCallback,
     token: CancellationToken,
@@ -364,15 +382,20 @@ export class LocalQueries extends DisposableObject {
     if (queryUri !== undefined) {
       // The query URI is provided by the command, most likely because the command was run from an
       // editor context menu. Use the provided URI, but make sure it's a valid query.
-      queryPath = validateQueryUri(queryUri, quickEval);
+      queryPath = validateQueryUri(queryUri, quickEval !== QuickEvalType.None);
     } else {
       // Use the currently selected query.
-      queryPath = await this.getCurrentQuery(quickEval);
+      queryPath = await this.getCurrentQuery(quickEval !== QuickEvalType.None);
     }
 
     const selectedQuery: SelectedQuery = {
       queryPath,
-      quickEval: quickEval ? await getQuickEvalContext(range) : undefined,
+      quickEval: quickEval
+        ? await getQuickEvalContext(
+            range,
+            quickEval === QuickEvalType.QuickEvalCount,
+          )
+        : undefined,
     };
 
     // If no databaseItem is specified, use the database currently selected in the Databases UI
@@ -392,6 +415,7 @@ export class LocalQueries extends DisposableObject {
       {
         queryPath: selectedQuery.queryPath,
         quickEvalPosition: selectedQuery.quickEval?.quickEvalPosition,
+        quickEvalCountOnly: selectedQuery.quickEval?.quickEvalCount,
       },
       true,
       additionalPacks,
@@ -481,7 +505,7 @@ export class LocalQueries extends DisposableObject {
       for (const item of quickpick) {
         try {
           await this.compileAndRunQuery(
-            false,
+            QuickEvalType.None,
             uri,
             progress,
             token,

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

@@ -68,6 +68,14 @@ export interface CompilationTarget {
  */
 export interface QuickEvalOptions {
   quickEvalPos?: Position;
+  /**
+   * Whether to only count the number of results.
+   *
+   * This is only supported by the new query server
+   * but it isn't worth having a separate type and
+   * it is fine to have an ignored optional field.
+   */
+  countOnly?: boolean;
 }
 
 /**

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

@@ -16,6 +16,10 @@ export interface CoreQueryTarget {
    * `query`.
    */
   quickEvalPosition?: Position;
+  /**
+   * If this is quick eval, whether to only count the number of results.
+   */
+  quickEvalCountOnly?: boolean;
 }
 
 export interface CoreQueryResults {

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

@@ -36,7 +36,10 @@ export async function compileAndRunQueryAgainstDatabaseCore(
   const target =
     query.quickEvalPosition !== undefined
       ? {
-          quickEval: { quickEvalPos: query.quickEvalPosition },
+          quickEval: {
+            quickEvalPos: query.quickEvalPosition,
+            countOnly: query.quickEvalCountOnly,
+          },
         }
       : { query: {} };
 

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

@@ -433,6 +433,7 @@ export function validateQueryPath(
 export interface QuickEvalContext {
   quickEvalPosition: messages.Position;
   quickEvalText: string;
+  quickEvalCount: boolean;
 }
 
 /**
@@ -443,6 +444,7 @@ export interface QuickEvalContext {
  */
 export async function getQuickEvalContext(
   range: Range | undefined,
+  isCountOnly: boolean,
 ): Promise<QuickEvalContext> {
   const editor = window.activeTextEditor;
   if (editor === undefined) {
@@ -465,6 +467,7 @@ export async function getQuickEvalContext(
   return {
     quickEvalPosition,
     quickEvalText,
+    quickEvalCount: isCountOnly,
   };
 }
 

extensions/ql-vscode/test/vscode-tests/cli-integration/queries.test.ts

@@ -28,7 +28,7 @@ import {
   QueryRunner,
 } from "../../../src/query-server/query-runner";
 import { SELECT_QUERY_NAME } from "../../../src/language-support";
-import { LocalQueries } from "../../../src/local-queries";
+import { LocalQueries, QuickEvalType } from "../../../src/local-queries";
 import { QueryResultType } from "../../../src/pure/new-messages";
 import { createVSCodeCommandManager } from "../../../src/common/vscode/commands";
 import {
@@ -45,7 +45,7 @@ async function compileAndRunQuery(
   mode: DebugMode,
   appCommands: AppCommandManager,
   localQueries: LocalQueries,
-  quickEval: boolean,
+  quickEval: QuickEvalType,
   queryUri: Uri,
   progress: ProgressCallback,
   token: CancellationToken,
@@ -184,7 +184,7 @@ describeWithCodeQL()("Queries", () => {
         mode,
         appCommandManager,
         localQueries,
-        false,
+        QuickEvalType.None,
         Uri.file(queryUsingExtensionPath),
         progress,
         token,
@@ -218,7 +218,7 @@ describeWithCodeQL()("Queries", () => {
         mode,
         appCommandManager,
         localQueries,
-        false,
+        QuickEvalType.None,
         Uri.file(queryPath),
         progress,
         token,
@@ -238,7 +238,7 @@ describeWithCodeQL()("Queries", () => {
         mode,
         appCommandManager,
         localQueries,
-        false,
+        QuickEvalType.None,
         Uri.file(queryPath),
         progress,
         token,

extensions/ql-vscode/test/vscode-tests/minimal-workspace/determining-selected-query-test.ts

@@ -26,7 +26,7 @@ export function run() {
 
     it("should allow ql files to be quick-evaled", async () => {
       await showQlDocument("query.ql");
-      const q = await getQuickEvalContext(undefined);
+      const q = await getQuickEvalContext(undefined, false);
       expect(
         q.quickEvalPosition.fileName.endsWith(
           join("ql-vscode", "test", "data", "query.ql"),
@@ -36,7 +36,7 @@ export function run() {
 
     it("should allow qll files to be quick-evaled", async () => {
       await showQlDocument("library.qll");
-      const q = await getQuickEvalContext(undefined);
+      const q = await getQuickEvalContext(undefined, false);
       expect(
         q.quickEvalPosition.fileName.endsWith(
           join("ql-vscode", "test", "data", "library.qll"),
@@ -55,7 +55,7 @@ export function run() {
 
     it("should reject non-ql[l] files when running a quick eval", async () => {
       await showQlDocument("textfile.txt");
-      await expect(getQuickEvalContext(undefined)).rejects.toThrow(
+      await expect(getQuickEvalContext(undefined, false)).rejects.toThrow(
         "The selected resource is not a CodeQL file",
       );
     });

extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts

@@ -77,6 +77,7 @@ describe("runQuery", () => {
         {
           queryPath: expect.stringMatching(/FetchExternalApis\.ql/),
           quickEvalPosition: undefined,
+          quickEvalCountOnly: false,
         },
         false,
         [],