Extract query resolving to more generic functions

koesie10 · koesie10 · commit 32656c1cb81f · 2023-07-25T10:03:01.000+02:00
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
@@ -14,7 +14,7 @@ import { QueryResultType } from "../query-server/new-messages";
 import { file } from "tmp-promise";
 import { writeFile } from "fs-extra";
 import { dump } from "js-yaml";
-import { qlpackOfDatabase } from "../language-support";
+import { qlpackOfDatabase } from "../local-queries";
 import { telemetryListener } from "../common/vscode/telemetry";
 
 type FlowModelOptions = {
diff --git a/extensions/ql-vscode/src/language-support/contextual/location-finder.ts b/extensions/ql-vscode/src/language-support/contextual/location-finder.ts
@@ -12,16 +12,13 @@ import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { DatabaseManager, DatabaseItem } from "../../databases/local-databases";
 import { ProgressCallback } from "../../common/vscode/progress";
 import { KeyType } from "./key-type";
-import {
-  qlpackOfDatabase,
-  resolveQueries,
-  runContextualQuery,
-} from "./query-resolver";
+import { resolveQueries, runContextualQuery } from "./query-resolver";
 import { CancellationToken, LocationLink, Uri } from "vscode";
 import { QueryOutputDir } from "../../run-queries-shared";
 import { QueryRunner } from "../../query-server";
 import { QueryResultType } from "../../query-server/new-messages";
 import { fileRangeFromURI } from "./file-range-from-uri";
+import { qlpackOfDatabase } from "../../local-queries";
 
 export const SELECT_QUERY_NAME = "#select";
 export const SELECTED_SOURCE_FILE = "selectedSourceFile";
diff --git a/extensions/ql-vscode/src/language-support/contextual/query-resolver.ts b/extensions/ql-vscode/src/language-support/contextual/query-resolver.ts
@@ -1,14 +1,8 @@
-import { writeFile, promises } from "fs-extra";
-import { dump } from "js-yaml";
-import { file } from "tmp-promise";
+import { promises } from "fs-extra";
 import { basename, dirname, resolve } from "path";
 
 import { getOnDiskWorkspaceFolders } from "../../common/vscode/workspace-folders";
-import {
-  getPrimaryDbscheme,
-  getQlPackForDbscheme,
-  QlPacksForLanguage,
-} from "../../databases/qlpack";
+import { QlPacksForLanguage } from "../../databases/qlpack";
 import {
   KeyType,
   kindOfKeyType,
@@ -17,97 +11,23 @@ import {
 } from "./key-type";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { DatabaseItem } from "../../databases/local-databases";
+import { resolveQueries as resolveLocalQueries } from "../../local-queries/query-resolver";
 import { extLogger } from "../../common/logging/vscode";
-import {
-  showAndLogExceptionWithTelemetry,
-  TeeLogger,
-} from "../../common/logging";
+import { TeeLogger } from "../../common/logging";
 import { CancellationToken } from "vscode";
 import { ProgressCallback } from "../../common/vscode/progress";
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
-import { redactableError } from "../../common/errors";
 import { QLPACK_FILENAMES } from "../../common/ql";
-import { telemetryListener } from "../../common/vscode/telemetry";
-
-export async function qlpackOfDatabase(
-  cli: Pick<CodeQLCliServer, "resolveQlpacks">,
-  db: Pick<DatabaseItem, "contents">,
-): Promise<QlPacksForLanguage> {
-  if (db.contents === undefined) {
-    throw new Error("Database is invalid and cannot infer QLPack.");
-  }
-  const datasetPath = db.contents.datasetUri.fsPath;
-  const dbscheme = await getPrimaryDbscheme(datasetPath);
-  return await getQlPackForDbscheme(cli, dbscheme);
-}
-
-/**
- * Finds the contextual queries with the specified key in a list of CodeQL packs.
- *
- * @param cli The CLI instance to use.
- * @param qlpacks The list of packs to search.
- * @param keyType The contextual query key of the query to search for.
- * @returns The found queries from the first pack in which any matching queries were found.
- */
-async function resolveQueriesFromPacks(
-  cli: CodeQLCliServer,
-  qlpacks: string[],
-  keyType: KeyType,
-): Promise<string[]> {
-  const suiteFile = (
-    await file({
-      postfix: ".qls",
-    })
-  ).path;
-  const suiteYaml = [];
-  for (const qlpack of qlpacks) {
-    suiteYaml.push({
-      from: qlpack,
-      queries: ".",
-      include: {
-        kind: kindOfKeyType(keyType),
-        "tags contain": tagOfKeyType(keyType),
-      },
-    });
-  }
-  await writeFile(suiteFile, dump(suiteYaml), "utf8");
-
-  const queries = await cli.resolveQueriesInSuite(
-    suiteFile,
-    getOnDiskWorkspaceFolders(),
-  );
-  return queries;
-}
 
 export async function resolveQueries(
   cli: CodeQLCliServer,
   qlpacks: QlPacksForLanguage,
   keyType: KeyType,
 ): Promise<string[]> {
-  const packsToSearch: string[] = [];
-
-  // The CLI can handle both library packs and query packs, so search both packs in order.
-  packsToSearch.push(qlpacks.dbschemePack);
-  if (qlpacks.queryPack !== undefined) {
-    packsToSearch.push(qlpacks.queryPack);
-  }
-
-  const queries = await resolveQueriesFromPacks(cli, packsToSearch, keyType);
-  if (queries.length > 0) {
-    return queries;
-  }
-
-  // No queries found. Determine the correct error message for the various scenarios.
-  const keyTypeName = nameOfKeyType(keyType);
-  const keyTypeTag = tagOfKeyType(keyType);
-  const joinedPacksToSearch = packsToSearch.join(", ");
-  const error = redactableError`No ${keyTypeName} queries (tagged "${keyTypeTag}") could be found in the \
-current library path (tried searching the following packs: ${joinedPacksToSearch}). \
-Try upgrading the CodeQL libraries. If that doesn't work, then ${keyTypeName} queries are not yet available \
-for this language.`;
-
-  void showAndLogExceptionWithTelemetry(extLogger, telemetryListener, error);
-  throw error;
+  return resolveLocalQueries(cli, qlpacks, nameOfKeyType(keyType), {
+    kind: kindOfKeyType(keyType),
+    "tags contain": [tagOfKeyType(keyType)],
+  });
 }
 
 async function resolveContextualQuery(
diff --git a/extensions/ql-vscode/src/language-support/contextual/template-provider.ts b/extensions/ql-vscode/src/language-support/contextual/template-provider.ts
@@ -27,18 +27,15 @@ import {
   SELECTED_SOURCE_LINE,
   SELECTED_SOURCE_COLUMN,
 } from "./location-finder";
-import {
-  qlpackOfDatabase,
-  resolveQueries,
-  runContextualQuery,
-} from "./query-resolver";
+import { resolveQueries, runContextualQuery } from "./query-resolver";
 import {
   isCanary,
   NO_CACHE_AST_VIEWER,
   NO_CACHE_CONTEXTUAL_QUERIES,
 } from "../../config";
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
 import { AstBuilder } from "../ast-viewer/ast-builder";
+import { qlpackOfDatabase } from "../../local-queries";
 
 /**
  * Runs templated CodeQL queries to find definitions in
diff --git a/extensions/ql-vscode/src/local-queries/index.ts b/extensions/ql-vscode/src/local-queries/index.ts
@@ -1,5 +1,6 @@
 export * from "./local-queries";
 export * from "./local-query-run";
+export * from "./query-resolver";
 export * from "./quick-eval-code-lens-provider";
 export * from "./quick-query";
 export * from "./results-view";
diff --git a/extensions/ql-vscode/src/local-queries/query-resolver.ts b/extensions/ql-vscode/src/local-queries/query-resolver.ts
@@ -0,0 +1,131 @@
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { DatabaseItem } from "../databases/local-databases";
+import {
+  getPrimaryDbscheme,
+  getQlPackForDbscheme,
+  QlPacksForLanguage,
+} from "../databases/qlpack";
+import { file } from "tmp-promise";
+import { writeFile } from "fs-extra";
+import { dump } from "js-yaml";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { redactableError } from "../common/errors";
+import { showAndLogExceptionWithTelemetry } from "../common/logging";
+import { extLogger } from "../common/logging/vscode";
+import { telemetryListener } from "../common/vscode/telemetry";
+
+export async function qlpackOfDatabase(
+  cli: Pick<CodeQLCliServer, "resolveQlpacks">,
+  db: Pick<DatabaseItem, "contents">,
+): Promise<QlPacksForLanguage> {
+  if (db.contents === undefined) {
+    throw new Error("Database is invalid and cannot infer QLPack.");
+  }
+  const datasetPath = db.contents.datasetUri.fsPath;
+  const dbscheme = await getPrimaryDbscheme(datasetPath);
+  return await getQlPackForDbscheme(cli, dbscheme);
+}
+
+export interface QueryConstraints {
+  kind?: string;
+  "tags contain"?: string[];
+  "tags contain all"?: string[];
+}
+
+/**
+ * Finds the queries with the specified kind and tags in a list of CodeQL packs.
+ *
+ * @param cli The CLI instance to use.
+ * @param qlpacks The list of packs to search.
+ * @param constraints Constraints on the queries to search for.
+ * @returns The found queries from the first pack in which any matching queries were found.
+ */
+async function resolveQueriesFromPacks(
+  cli: CodeQLCliServer,
+  qlpacks: string[],
+  constraints: QueryConstraints,
+): Promise<string[]> {
+  const suiteFile = (
+    await file({
+      postfix: ".qls",
+    })
+  ).path;
+  const suiteYaml = [];
+  for (const qlpack of qlpacks) {
+    suiteYaml.push({
+      from: qlpack,
+      queries: ".",
+      include: constraints,
+    });
+  }
+  await writeFile(
+    suiteFile,
+    dump(suiteYaml, {
+      noRefs: true, // CodeQL doesn't really support refs
+    }),
+    "utf8",
+  );
+
+  return await cli.resolveQueriesInSuite(
+    suiteFile,
+    getOnDiskWorkspaceFolders(),
+  );
+}
+
+/**
+ * Finds the queries with the specified kind and tags in a QLPack.
+ *
+ * @param cli The CLI instance to use.
+ * @param qlpacks The list of packs to search.
+ * @param name The name of the query to use in error messages.
+ * @param constraints Constraints on the queries to search for.
+ * @returns The found queries from the first pack in which any matching queries were found.
+ */
+export async function resolveQueries(
+  cli: CodeQLCliServer,
+  qlpacks: QlPacksForLanguage,
+  name: string,
+  constraints: QueryConstraints,
+): Promise<string[]> {
+  const packsToSearch: string[] = [];
+
+  // The CLI can handle both library packs and query packs, so search both packs in order.
+  packsToSearch.push(qlpacks.dbschemePack);
+  if (qlpacks.queryPack !== undefined) {
+    packsToSearch.push(qlpacks.queryPack);
+  }
+
+  const queries = await resolveQueriesFromPacks(
+    cli,
+    packsToSearch,
+    constraints,
+  );
+  if (queries.length > 0) {
+    return queries;
+  }
+
+  // No queries found. Determine the correct error message for the various scenarios.
+  const humanConstraints = [];
+  if (constraints.kind !== undefined) {
+    humanConstraints.push(`kind "${constraints.kind}"`);
+  }
+  if (constraints["tags contain"] !== undefined) {
+    humanConstraints.push(`tagged "${constraints["tags contain"].join(" ")}"`);
+  }
+  if (constraints["tags contain all"] !== undefined) {
+    humanConstraints.push(
+      `tagged all of "${constraints["tags contain all"].join(" ")}"`,
+    );
+  }
+
+  const joinedPacksToSearch = packsToSearch.join(", ");
+  const error = redactableError`No ${name} queries (${humanConstraints.join(
+    ", ",
+  )}) could be found in the \
+current library path (tried searching the following packs: ${joinedPacksToSearch}). \
+Try upgrading the CodeQL libraries. If that doesn't work, then ${name} queries are not yet available \
+for this language.`;
+
+  void showAndLogExceptionWithTelemetry(extLogger, telemetryListener, error);
+  throw error;
+}
diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/query-resolver.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/query-resolver.test.ts
@@ -5,40 +5,18 @@ import { getErrorMessage } from "../../../../../src/common/helpers-pure";
 
 import * as log from "../../../../../src/common/logging/notifications";
 import * as workspaceFolders from "../../../../../src/common/vscode/workspace-folders";
-import * as qlpack from "../../../../../src/databases/qlpack";
-import {
-  KeyType,
-  qlpackOfDatabase,
-  resolveQueries,
-} from "../../../../../src/language-support";
+import { KeyType, resolveQueries } from "../../../../../src/language-support";
 import { CodeQLCliServer } from "../../../../../src/codeql-cli/cli";
-import { mockDatabaseItem, mockedObject } from "../../../utils/mocking.helpers";
+import { mockedObject } from "../../../utils/mocking.helpers";
 
 describe("queryResolver", () => {
-  let getQlPackForDbschemeSpy: jest.SpiedFunction<
-    typeof qlpack.getQlPackForDbscheme
-  >;
-  let getPrimaryDbschemeSpy: jest.SpiedFunction<
-    typeof qlpack.getPrimaryDbscheme
-  >;
-
   const resolveQueriesInSuite = jest.fn();
 
   const mockCli = mockedObject<CodeQLCliServer>({
     resolveQueriesInSuite,
   });
 
   beforeEach(() => {
-    getQlPackForDbschemeSpy = jest
-      .spyOn(qlpack, "getQlPackForDbscheme")
-      .mockResolvedValue({
-        dbschemePack: "dbschemePack",
-        dbschemePackIsLibraryPack: false,
-      });
-    getPrimaryDbschemeSpy = jest
-      .spyOn(qlpack, "getPrimaryDbscheme")
-      .mockResolvedValue("primaryDbscheme");
-
     jest
       .spyOn(workspaceFolders, "getOnDiskWorkspaceFolders")
       .mockReturnValue([]);
@@ -68,7 +46,7 @@ describe("queryResolver", () => {
           queries: ".",
           include: {
             kind: "definitions",
-            "tags contain": "ide-contextual-queries/local-definitions",
+            "tags contain": ["ide-contextual-queries/local-definitions"],
           },
         },
       ]);
@@ -87,31 +65,9 @@ describe("queryResolver", () => {
         expect(true).toBe(false);
       } catch (e) {
         expect(getErrorMessage(e)).toBe(
-          'No definitions queries (tagged "ide-contextual-queries/local-definitions") could be found in the current library path (tried searching the following packs: my-qlpack). Try upgrading the CodeQL libraries. If that doesn\'t work, then definitions queries are not yet available for this language.',
+          'No definitions queries (kind "definitions", tagged "ide-contextual-queries/local-definitions") could be found in the current library path (tried searching the following packs: my-qlpack). Try upgrading the CodeQL libraries. If that doesn\'t work, then definitions queries are not yet available for this language.',
         );
       }
     });
   });
-
-  describe("qlpackOfDatabase", () => {
-    it("should get the qlpack of a database", async () => {
-      getQlPackForDbschemeSpy.mockResolvedValue({
-        dbschemePack: "my-qlpack",
-        dbschemePackIsLibraryPack: false,
-      });
-      const db = mockDatabaseItem({
-        contents: {
-          datasetUri: {
-            fsPath: "/path/to/database",
-          },
-        },
-      });
-      const result = await qlpackOfDatabase(mockCli, db);
-      expect(result).toEqual({
-        dbschemePack: "my-qlpack",
-        dbschemePackIsLibraryPack: false,
-      });
-      expect(getPrimaryDbschemeSpy).toBeCalledWith("/path/to/database");
-    });
-  });
 });
diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/local-queries/query-resolver.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/local-queries/query-resolver.test.ts
