From 9cd6dafdf4f50b2eb23f91aa7ce0d806f252c6b6 Mon Sep 17 00:00:00 2001 From: Shati Patel <42641846+shati-patel@users.noreply.github.com> Date: Wed, 31 Jan 2024 13:04:08 +0000 Subject: [PATCH] Functions for running "access path suggestions" queries (#3294) --- .../src/model-editor/suggestion-queries.ts | 180 +++++++++++++++ .../ql-vscode/src/model-editor/suggestions.ts | 5 + .../model-editor/suggestion-queries.test.ts | 216 ++++++++++++++++++ 3 files changed, 401 insertions(+) create mode 100644 extensions/ql-vscode/src/model-editor/suggestion-queries.ts create mode 100644 extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts diff --git a/extensions/ql-vscode/src/model-editor/suggestion-queries.ts b/extensions/ql-vscode/src/model-editor/suggestion-queries.ts new file mode 100644 index 000000000..ef96da2d3 --- /dev/null +++ b/extensions/ql-vscode/src/model-editor/suggestion-queries.ts @@ -0,0 +1,180 @@ +import type { CodeQLCliServer } from "../codeql-cli/cli"; +import type { Mode } from "./shared/mode"; +import { resolveQueriesFromPacks } from "../local-queries"; +import { modeTag } from "./mode-tag"; +import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders"; +import type { NotificationLogger } from "../common/logging"; +import { showAndLogExceptionWithTelemetry } from "../common/logging"; +import { telemetryListener } from "../common/vscode/telemetry"; +import { redactableError } from "../common/errors"; +import { runQuery } from "../local-queries/run-query"; +import type { QueryRunner } from "../query-server"; +import type { DatabaseItem } from "../databases/local-databases"; +import type { ProgressCallback } from "../common/vscode/progress"; +import type { CancellationToken } from "vscode"; +import type { DecodedBqrsChunk } from "../common/bqrs-cli-types"; +import type { + AccessPathSuggestionRow, + AccessPathSuggestionRows, +} from "./suggestions"; + +type RunQueryOptions = { + parseResults: ( + results: DecodedBqrsChunk, + ) => AccessPathSuggestionRow[] | Promise; + + cliServer: CodeQLCliServer; + queryRunner: QueryRunner; + logger: NotificationLogger; + databaseItem: DatabaseItem; + queryStorageDir: string; + + progress: ProgressCallback; + token: CancellationToken; +}; + +const maxStep = 2000; + +export async function runSuggestionsQuery( + mode: Mode, + { + parseResults, + cliServer, + queryRunner, + logger, + databaseItem, + queryStorageDir, + progress, + token, + }: RunQueryOptions, +): Promise { + progress({ + message: "Resolving QL packs", + step: 1, + maxStep, + }); + const additionalPacks = getOnDiskWorkspaceFolders(); + const extensionPacks = Object.keys( + await cliServer.resolveQlpacks(additionalPacks, true), + ); + + progress({ + message: "Resolving query", + step: 2, + maxStep, + }); + + const queryPath = await resolveSuggestionsQuery( + cliServer, + databaseItem.language, + mode, + ); + if (!queryPath) { + void showAndLogExceptionWithTelemetry( + logger, + telemetryListener, + redactableError`The ${mode} access path suggestions query could not be found. Try re-opening the model editor. If that doesn't work, try upgrading the CodeQL libraries.`, + ); + return undefined; + } + + // Run the actual query + const completedQuery = await runQuery({ + queryRunner, + databaseItem, + queryPath, + queryStorageDir, + additionalPacks, + extensionPacks, + progress: (update) => + progress({ + step: update.step + 500, + maxStep, + message: update.message, + }), + token, + }); + + if (!completedQuery) { + return undefined; + } + + // Read the results and convert to internal representation + progress({ + message: "Decoding results", + step: 1600, + maxStep, + }); + + const bqrs = await cliServer.bqrsDecodeAll(completedQuery.outputDir.bqrsPath); + + progress({ + message: "Finalizing results", + step: 1950, + maxStep, + }); + + const inputChunk = bqrs["input"]; + const outputChunk = bqrs["output"]; + + if (!inputChunk && !outputChunk) { + void logger.log( + `No results found for ${mode} access path suggestions query`, + ); + return undefined; + } + + const inputSuggestions = inputChunk ? await parseResults(inputChunk) : []; + const outputSuggestions = outputChunk ? await parseResults(outputChunk) : []; + + return { + input: inputSuggestions, + output: outputSuggestions, + }; +} + +/** + * Resolve the query path to the model editor access path suggestions query. All queries are tagged like this: + * modeleditor access-path-suggestions + * Example: modeleditor access-path-suggestions framework-mode + * + * @param cliServer The CodeQL CLI server to use. + * @param language The language of the query pack to use. + * @param mode The mode to resolve the query for. + * @param additionalPackNames Additional pack names to search. + * @param additionalPackPaths Additional pack paths to search. + */ +async function resolveSuggestionsQuery( + cliServer: CodeQLCliServer, + language: string, + mode: Mode, + additionalPackNames: string[] = [], + additionalPackPaths: string[] = [], +): Promise { + const packsToSearch = [`codeql/${language}-queries`, ...additionalPackNames]; + + const queries = await resolveQueriesFromPacks( + cliServer, + packsToSearch, + { + kind: "table", + "tags contain all": [ + "modeleditor", + "access-path-suggestions", + modeTag(mode), + ], + }, + additionalPackPaths, + ); + if (queries.length > 1) { + throw new Error( + `Found multiple suggestions queries for ${mode}. Can't continue`, + ); + } + + if (queries.length === 0) { + return undefined; + } + + return queries[0]; +} diff --git a/extensions/ql-vscode/src/model-editor/suggestions.ts b/extensions/ql-vscode/src/model-editor/suggestions.ts index 0cc7c04ad..cc3211c26 100644 --- a/extensions/ql-vscode/src/model-editor/suggestions.ts +++ b/extensions/ql-vscode/src/model-editor/suggestions.ts @@ -25,6 +25,11 @@ export type AccessPathSuggestionRow = { details: string; }; +export type AccessPathSuggestionRows = { + input: AccessPathSuggestionRow[]; + output: AccessPathSuggestionRow[]; +}; + export type AccessPathOption = { label: string; value: string; diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts new file mode 100644 index 000000000..336acb2fd --- /dev/null +++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts @@ -0,0 +1,216 @@ +import { createMockLogger } from "../../../__mocks__/loggerMock"; +import type { DatabaseItem } from "../../../../src/databases/local-databases"; +import { DatabaseKind } from "../../../../src/databases/local-databases"; +import { file } from "tmp-promise"; +import { QueryResultType } from "../../../../src/query-server/messages"; +import { QueryLanguage } from "../../../../src/common/query-language"; +import { mockedObject, mockedUri } from "../../utils/mocking.helpers"; +import { Mode } from "../../../../src/model-editor/shared/mode"; +import { join } from "path"; +import type { CodeQLCliServer } from "../../../../src/codeql-cli/cli"; +import type { QueryRunner } from "../../../../src/query-server"; +import { QueryOutputDir } from "../../../../src/local-queries/query-output-dir"; +import { runSuggestionsQuery } from "../../../../src/model-editor/suggestion-queries"; + +describe("runSuggestionsQuery", () => { + const mockDecodedBqrs = { + input: { + columns: [ + { + name: "type", + kind: "String", + }, + { + name: "path", + kind: "String", + }, + { + name: "value", + kind: "String", + }, + { + name: "details", + kind: "String", + }, + { + name: "defType", + kind: "String", + }, + ], + tuples: [ + [ + "Correctness", + "Method[assert!]", + "Argument[self]", + "self in assert!", + "parameter", + ], + ], + }, + output: { + columns: [ + { + name: "type", + kind: "String", + }, + { + name: "path", + kind: "String", + }, + { + name: "value", + kind: "String", + }, + { + name: "details", + kind: "String", + }, + { + name: "defType", + kind: "String", + }, + ], + tuples: [ + [ + "Correctness", + "Method[assert!]", + "ReturnValue", + "call to puts", + "return", + ], + [ + "Correctness", + "Method[assert!]", + "Argument[self]", + "self in assert!", + "parameter", + ], + ], + }, + }; + const mockInputSuggestions = [ + { + method: { + packageName: "", + typeName: "Correctness", + methodName: "assert!", + methodParameters: "", + signature: "Correctness#assert!", + }, + value: "Argument[self]", + details: "self in assert!", + definitionType: "parameter", + }, + ]; + const mockOutputSuggestions = [ + { + method: { + packageName: "", + typeName: "Correctness", + methodName: "assert!", + methodParameters: "", + signature: "Correctness#assert!", + }, + value: "ReturnValue", + details: "call to puts", + definitionType: "return", + }, + { + method: { + packageName: "", + typeName: "Correctness", + methodName: "assert!", + methodParameters: "", + signature: "Correctness#assert!", + }, + value: "Argument[self]", + details: "self in assert!", + definitionType: "parameter", + }, + ]; + + it("should run query", async () => { + const language = QueryLanguage.Ruby; + const outputDir = new QueryOutputDir(join((await file()).path, "1")); + + const parseResults = jest + .fn() + .mockResolvedValueOnce(mockInputSuggestions) + .mockResolvedValueOnce(mockOutputSuggestions); + + const options = { + parseResults, + cliServer: mockedObject({ + resolveQlpacks: jest.fn().mockResolvedValue({ + "my/extensions": "/a/b/c/", + }), + resolveQueriesInSuite: jest + .fn() + .mockResolvedValue(["/a/b/c/FrameworkModeAccessPathSuggestions.ql"]), + packPacklist: jest + .fn() + .mockResolvedValue([ + "/a/b/c/qlpack.yml", + "/a/b/c/qlpack.lock.yml", + "/a/b/c/qlpack2.yml", + ]), + bqrsDecodeAll: jest.fn().mockResolvedValue(mockDecodedBqrs), + }), + queryRunner: mockedObject({ + createQueryRun: jest.fn().mockReturnValue({ + evaluate: jest.fn().mockResolvedValue({ + resultType: QueryResultType.SUCCESS, + outputDir, + }), + outputDir, + }), + logger: createMockLogger(), + }), + logger: createMockLogger(), + databaseItem: mockedObject({ + databaseUri: mockedUri("/a/b/c/src.zip"), + contents: { + kind: DatabaseKind.Database, + name: "foo", + datasetUri: mockedUri(), + }, + language, + }), + queryStorageDir: "/tmp/queries", + progress: jest.fn(), + token: { + isCancellationRequested: false, + onCancellationRequested: jest.fn(), + }, + }; + + const result = await runSuggestionsQuery(Mode.Framework, options); + + expect(result).not.toBeUndefined; + + expect(options.cliServer.resolveQlpacks).toHaveBeenCalledTimes(1); + expect(options.cliServer.resolveQlpacks).toHaveBeenCalledWith([], true); + expect(options.queryRunner.createQueryRun).toHaveBeenCalledWith( + "/a/b/c/src.zip", + { + queryPath: expect.stringMatching(/\S*AccessPathSuggestions\.ql/), + quickEvalPosition: undefined, + quickEvalCountOnly: false, + }, + false, + [], + ["my/extensions"], + {}, + "/tmp/queries", + undefined, + undefined, + ); + + expect(options.parseResults).toHaveBeenCalledTimes(2); + + expect(result).toEqual({ + input: mockInputSuggestions, + output: mockOutputSuggestions, + }); + }); +});