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

Author: Alexander Eyers-Taylor

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: {} };