From 22172d5d74e5633dd65c47fb04ea3da0183e7140 Mon Sep 17 00:00:00 2001 From: Alexander Eyers-Taylor Date: Tue, 30 May 2023 14:12:47 +0100 Subject: [PATCH] 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 --- .../external-api-usage-query.ts | 6 ++- .../generate-flow-model.ts | 6 ++- .../src/debugger/debug-configuration.ts | 2 +- .../ql-vscode/src/debugger/debug-session.ts | 1 + .../ql-vscode/src/debugger/debugger-ui.ts | 2 +- .../ast-viewer/ast-cfg-commands.ts | 4 +- .../src/local-queries/local-queries.ts | 44 ++++++++++++++----- .../ql-vscode/src/pure/messages-shared.ts | 8 ++++ .../src/query-server/query-runner.ts | 4 ++ .../ql-vscode/src/query-server/run-queries.ts | 5 ++- .../ql-vscode/src/run-queries-shared.ts | 3 ++ .../cli-integration/queries.test.ts | 10 ++--- .../determining-selected-query-test.ts | 6 +-- .../external-api-usage-query.test.ts | 1 + 14 files changed, 77 insertions(+), 25 deletions(-) diff --git a/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts b/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts index 9fbc7b293..659f8d519 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/external-api-usage-query.ts +++ b/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, diff --git a/extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts b/extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts index 6f48d8587..d8330953e 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts +++ b/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, diff --git a/extensions/ql-vscode/src/debugger/debug-configuration.ts b/extensions/ql-vscode/src/debugger/debug-configuration.ts index c25e4fc6f..4d4a788a7 100644 --- a/extensions/ql-vscode/src/debugger/debug-configuration.ts +++ b/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 = { diff --git a/extensions/ql-vscode/src/debugger/debug-session.ts b/extensions/ql-vscode/src/debugger/debug-session.ts index 9b9091092..02c183200 100644 --- a/extensions/ql-vscode/src/debugger/debug-session.ts +++ b/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, diff --git a/extensions/ql-vscode/src/debugger/debugger-ui.ts b/extensions/ql-vscode/src/debugger/debugger-ui.ts index e6a115a43..46168e021 100644 --- a/extensions/ql-vscode/src/debugger/debugger-ui.ts +++ b/extensions/ql-vscode/src/debugger/debugger-ui.ts @@ -74,7 +74,7 @@ class QLDebugAdapterTracker public async quickEval(): Promise { const args: CodeQLProtocol.QuickEvalRequest["arguments"] = { - quickEvalContext: await getQuickEvalContext(undefined), + quickEvalContext: await getQuickEvalContext(undefined, false), }; await this.session.customRequest("codeql-quickeval", args); } diff --git a/extensions/ql-vscode/src/language-support/ast-viewer/ast-cfg-commands.ts b/extensions/ql-vscode/src/language-support/ast-viewer/ast-cfg-commands.ts index 3325c9df0..3262fbc8f 100644 --- a/extensions/ql-vscode/src/language-support/ast-viewer/ast-cfg-commands.ts +++ b/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, diff --git a/extensions/ql-vscode/src/local-queries/local-queries.ts b/extensions/ql-vscode/src/local-queries/local-queries.ts index b7250cbb4..a5b7f3028 100644 --- a/extensions/ql-vscode/src/local-queries/local-queries.ts +++ b/extensions/ql-vscode/src/local-queries/local-queries.ts @@ -72,6 +72,12 @@ async function promptToSaveQueryIfNeeded(query: SelectedQuery): Promise { } } +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 { 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 { 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, diff --git a/extensions/ql-vscode/src/pure/messages-shared.ts b/extensions/ql-vscode/src/pure/messages-shared.ts index 23b5a6871..a07d81f82 100644 --- a/extensions/ql-vscode/src/pure/messages-shared.ts +++ b/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; } /** diff --git a/extensions/ql-vscode/src/query-server/query-runner.ts b/extensions/ql-vscode/src/query-server/query-runner.ts index 95d7eafef..03f886ffd 100644 --- a/extensions/ql-vscode/src/query-server/query-runner.ts +++ b/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 { diff --git a/extensions/ql-vscode/src/query-server/run-queries.ts b/extensions/ql-vscode/src/query-server/run-queries.ts index 9f9e441c6..66674bab9 100644 --- a/extensions/ql-vscode/src/query-server/run-queries.ts +++ b/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: {} }; diff --git a/extensions/ql-vscode/src/run-queries-shared.ts b/extensions/ql-vscode/src/run-queries-shared.ts index cdb472a0b..811fe57df 100644 --- a/extensions/ql-vscode/src/run-queries-shared.ts +++ b/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 { const editor = window.activeTextEditor; if (editor === undefined) { @@ -465,6 +467,7 @@ export async function getQuickEvalContext( return { quickEvalPosition, quickEvalText, + quickEvalCount: isCountOnly, }; } diff --git a/extensions/ql-vscode/test/vscode-tests/cli-integration/queries.test.ts b/extensions/ql-vscode/test/vscode-tests/cli-integration/queries.test.ts index 491e6d037..d068a1295 100644 --- a/extensions/ql-vscode/test/vscode-tests/cli-integration/queries.test.ts +++ b/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, diff --git a/extensions/ql-vscode/test/vscode-tests/minimal-workspace/determining-selected-query-test.ts b/extensions/ql-vscode/test/vscode-tests/minimal-workspace/determining-selected-query-test.ts index 921bf037f..61eeb724f 100644 --- a/extensions/ql-vscode/test/vscode-tests/minimal-workspace/determining-selected-query-test.ts +++ b/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", ); }); diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts index a4b5ec0e0..548637cc6 100644 --- a/extensions/ql-vscode/test/vscode-tests/no-workspace/data-extensions-editor/external-api-usage-query.test.ts +++ b/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, [],