Functions for running "access path suggestions" queries (#3294)

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<AccessPathSuggestionRow[]>;
+
+  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<AccessPathSuggestionRows | undefined> {
+  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 <mode>
+ * 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<string | undefined> {
+  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];
+}

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;

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<CodeQLCliServer>({
+        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<QueryRunner>({
+        createQueryRun: jest.fn().mockReturnValue({
+          evaluate: jest.fn().mockResolvedValue({
+            resultType: QueryResultType.SUCCESS,
+            outputDir,
+          }),
+          outputDir,
+        }),
+        logger: createMockLogger(),
+      }),
+      logger: createMockLogger(),
+      databaseItem: mockedObject<DatabaseItem>({
+        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,
+    });
+  });
+});