Make query resolver more generic

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 = {

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";

extensions/ql-vscode/src/language-support/contextual/query-resolver.ts

@@ -1,14 +1,5 @@
-import { writeFile, promises } from "fs-extra";
-import { dump } from "js-yaml";
-import { file } from "tmp-promise";
-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,154 +8,22 @@ 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;
-}
+import { createLockFileForStandardQuery } from "../../local-queries/standard-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;
-}
-
-async function resolveContextualQuery(
-  cli: CodeQLCliServer,
-  query: string,
-): Promise<{ packPath: string; createdTempLockFile: boolean }> {
-  // Contextual queries now live within the standard library packs.
-  // This simplifies distribution (you don't need the standard query pack to use the AST viewer),
-  // but if the library pack doesn't have a lockfile, we won't be able to find
-  // other pack dependencies of the library pack.
-
-  // Work out the enclosing pack.
-  const packContents = await cli.packPacklist(query, false);
-  const packFilePath = packContents.find((p) =>
-    QLPACK_FILENAMES.includes(basename(p)),
-  );
-  if (packFilePath === undefined) {
-    // Should not happen; we already resolved this query.
-    throw new Error(
-      `Could not find a CodeQL pack file for the pack enclosing the contextual query ${query}`,
-    );
-  }
-  const packPath = dirname(packFilePath);
-  const lockFilePath = packContents.find((p) =>
-    ["codeql-pack.lock.yml", "qlpack.lock.yml"].includes(basename(p)),
-  );
-  let createdTempLockFile = false;
-  if (!lockFilePath) {
-    // No lock file, likely because this library pack is in the package cache.
-    // Create a lock file so that we can resolve dependencies and library path
-    // for the contextual query.
-    void extLogger.log(
-      `Library pack ${packPath} is missing a lock file; creating a temporary lock file`,
-    );
-    await cli.packResolveDependencies(packPath);
-    createdTempLockFile = true;
-    // Clear CLI server pack cache before installing dependencies,
-    // so that it picks up the new lock file, not the previously cached pack.
-    void extLogger.log("Clearing the CodeQL CLI server's pack cache");
-    await cli.clearCache();
-    // Install dependencies.
-    void extLogger.log(
-      `Installing package dependencies for library pack ${packPath}`,
-    );
-    await cli.packInstall(packPath);
-  }
-  return { packPath, createdTempLockFile };
-}
-
-async function removeTemporaryLockFile(packPath: string) {
-  const tempLockFilePath = resolve(packPath, "codeql-pack.lock.yml");
-  void extLogger.log(
-    `Deleting temporary package lock file at ${tempLockFilePath}`,
-  );
-  // It's fine if the file doesn't exist.
-  await promises.rm(resolve(packPath, "codeql-pack.lock.yml"), {
-    force: true,
+  return resolveLocalQueries(cli, qlpacks, nameOfKeyType(keyType), {
+    kind: kindOfKeyType(keyType),
+    "tags contain": [tagOfKeyType(keyType)],
   });
 }
 
@@ -178,10 +37,7 @@ export async function runContextualQuery(
   token: CancellationToken,
   templates: Record<string, string>,
 ): Promise<CoreCompletedQuery> {
-  const { packPath, createdTempLockFile } = await resolveContextualQuery(
-    cli,
-    query,
-  );
+  const { cleanup } = await createLockFileForStandardQuery(cli, query);
   const queryRun = qs.createQueryRun(
     db.databaseUri.fsPath,
     { queryPath: query, quickEvalPosition: undefined },
@@ -200,8 +56,6 @@ export async function runContextualQuery(
     token,
     new TeeLogger(qs.logger, queryRun.outputDir.logPath),
   );
-  if (createdTempLockFile) {
-    await removeTemporaryLockFile(packPath);
-  }
+  await cleanup?.();
   return results;
 }

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

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";

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;
+}