Make contextual queries work for fresh installs

This fixes the contextual queries when you are not in a workspace with
the submodule and do not have any downloaded packs in the package cache.
In that case, the contextual queries would fail because they weren't
able to determine which pack belonged to the database.

This fixes it by downloading the `codeql/${language}-all` pack in case
no dbscheme is found for the database. After the download is complete,
it will return the expected value for the qlpacks. This should work in
almost all cases (at least for standard languages).

Author: koesie10

extensions/ql-vscode/src/language-support/contextual/location-finder.ts

@@ -9,16 +9,19 @@ import {
   ResultSetSchema,
 } from "../../common/bqrs-cli-types";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
-import { DatabaseManager, DatabaseItem } from "../../databases/local-databases";
+import { DatabaseItem, DatabaseManager } from "../../databases/local-databases";
 import { ProgressCallback } from "../../common/vscode/progress";
 import { KeyType } from "./key-type";
-import { resolveQueries, runContextualQuery } from "./query-resolver";
+import {
+  resolveContextualQlPacksForDatabase,
+  resolveContextualQueries,
+  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";
@@ -63,11 +66,11 @@ export async function getLocationsForUriString(
     return [];
   }
 
-  const qlpack = await qlpackOfDatabase(cli, db);
+  const qlpack = await resolveContextualQlPacksForDatabase(cli, db);
   const templates = createTemplates(uri.pathWithinSourceArchive);
 
   const links: FullLocationLink[] = [];
-  for (const query of await resolveQueries(cli, qlpack, keyType)) {
+  for (const query of await resolveContextualQueries(cli, qlpack, keyType)) {
     const results = await runContextualQuery(
       query,
       db,

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

@@ -8,23 +8,67 @@ import {
 } from "./key-type";
 import { CodeQLCliServer } from "../../codeql-cli/cli";
 import { DatabaseItem } from "../../databases/local-databases";
-import { resolveQueriesByLanguagePack as resolveLocalQueries } from "../../local-queries/query-resolver";
+import {
+  qlpackOfDatabase,
+  resolveQueriesByLanguagePack as resolveLocalQueriesByLanguagePack,
+} from "../../local-queries/query-resolver";
 import { extLogger } from "../../common/logging/vscode";
 import { TeeLogger } from "../../common/logging";
 import { CancellationToken } from "vscode";
 import { ProgressCallback } from "../../common/vscode/progress";
 import { CoreCompletedQuery, QueryRunner } from "../../query-server";
 import { createLockFileForStandardQuery } from "../../local-queries/standard-queries";
 
-export async function resolveQueries(
+/**
+ * This wil try to determine the qlpacks for a given database. If it can't find a matching
+ * dbscheme with downloaded packs, it will download the default packs instead.
+ *
+ * @param cli The CLI server to use
+ * @param databaseItem The database item to find the qlpacks for
+ */
+export async function resolveContextualQlPacksForDatabase(
+  cli: CodeQLCliServer,
+  databaseItem: DatabaseItem,
+): Promise<QlPacksForLanguage> {
+  try {
+    return await qlpackOfDatabase(cli, databaseItem);
+  } catch (e) {
+    // If we can't find the qlpacks for the database, use the defaults instead
+  }
+
+  const dbInfo = await cli.resolveDatabase(databaseItem.databaseUri.fsPath);
+  const primaryLanguage = dbInfo.languages?.[0];
+  if (!primaryLanguage) {
+    throw new Error("Unable to determine primary language of database");
+  }
+
+  const libraryPack = `codeql/${primaryLanguage}-all`;
+  const queryPack = `codeql/${primaryLanguage}-queries`;
+
+  await cli.packDownload([libraryPack, queryPack]);
+
+  // Return the default packs. If these weren't valid packs, the download would have failed.
+  return {
+    dbschemePack: libraryPack,
+    dbschemePackIsLibraryPack: true,
+    queryPack,
+  };
+}
+
+export async function resolveContextualQueries(
   cli: CodeQLCliServer,
   qlpacks: QlPacksForLanguage,
   keyType: KeyType,
 ): Promise<string[]> {
-  return resolveLocalQueries(cli, qlpacks, nameOfKeyType(keyType), {
-    kind: kindOfKeyType(keyType),
-    "tags contain": [tagOfKeyType(keyType)],
-  });
+  return resolveLocalQueriesByLanguagePack(
+    cli,
+    qlpacks,
+    nameOfKeyType(keyType),
+    {
+      kind: kindOfKeyType(keyType),
+      "tags contain": [tagOfKeyType(keyType)],
+    },
+  );
 }
 
 export async function runContextualQuery(

extensions/ql-vscode/src/language-support/contextual/template-provider.ts

@@ -23,19 +23,22 @@ import { KeyType } from "./key-type";
 import {
   FullLocationLink,
   getLocationsForUriString,
+  SELECTED_SOURCE_COLUMN,
   SELECTED_SOURCE_FILE,
   SELECTED_SOURCE_LINE,
-  SELECTED_SOURCE_COLUMN,
 } from "./location-finder";
-import { resolveQueries, runContextualQuery } from "./query-resolver";
+import {
+  resolveContextualQlPacksForDatabase,
+  resolveContextualQueries,
+  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";
 import { MultiCancellationToken } from "../../common/vscode/multi-cancellation-token";
 
 /**
@@ -248,8 +251,8 @@ export class TemplatePrintAstProvider {
       throw new Error("Can't infer database from the provided source.");
     }
 
-    const qlpacks = await qlpackOfDatabase(this.cli, db);
-    const queries = await resolveQueries(
+    const qlpacks = await resolveContextualQlPacksForDatabase(this.cli, db);
+    const queries = await resolveContextualQueries(
       this.cli,
       qlpacks,
       KeyType.PrintAstQuery,
@@ -336,11 +339,11 @@ export class TemplatePrintCfgProvider {
       throw new Error("Can't infer database from the provided source.");
     }
 
-    const qlpack = await qlpackOfDatabase(this.cli, db);
+    const qlpack = await resolveContextualQlPacksForDatabase(this.cli, db);
     if (!qlpack) {
       throw new Error("Can't infer qlpack from database source archive.");
     }
-    const queries = await resolveQueries(
+    const queries = await resolveContextualQueries(
       this.cli,
       qlpack,
       KeyType.PrintCfgQuery,

extensions/ql-vscode/src/local-queries/query-resolver.ts

@@ -16,6 +16,11 @@ import { telemetryListener } from "../common/vscode/telemetry";
 import { SuiteInstruction } from "../packaging/suite-instruction";
 import { QueryConstraints } from "./query-constraints";
 
+/**
+ * Consider using `resolveContextualQlPacksForDatabase` instead.
+ * @param cli The CLI server instance to use.
+ * @param db The database to find the QLPack for.
+ */
 export async function qlpackOfDatabase(
   cli: Pick<CodeQLCliServer, "resolveQlpacks">,
   db: Pick<DatabaseItem, "contents">,

extensions/ql-vscode/test/vscode-tests/cli-integration/run-cli.test.ts

@@ -8,7 +8,10 @@ import {
 } from "../../../src/codeql-cli/cli";
 import { itWithCodeQL } from "../cli";
 import { getOnDiskWorkspaceFolders } from "../../../src/common/vscode/workspace-folders";
-import { KeyType, resolveQueries } from "../../../src/language-support";
+import {
+  KeyType,
+  resolveContextualQueries,
+} from "../../../src/language-support";
 import { faker } from "@faker-js/faker";
 import { getActivatedExtension } from "../global.helper";
 import { BaseLogger } from "../../../src/common/logging";
@@ -117,7 +120,11 @@ describe("Use cli", () => {
           expect(pack.queryPack).toContain(lang);
         }
 
-        const result = await resolveQueries(cli, pack, KeyType.PrintAstQuery);
+        const result = await resolveContextualQueries(
+          cli,
+          pack,
+          KeyType.PrintAstQuery,
+        );
 
         // It doesn't matter what the name or path of the query is, only
         // that we have found exactly one query.

extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/query-resolver.test.ts

@@ -5,28 +5,97 @@ 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 { KeyType, resolveQueries } from "../../../../../src/language-support";
-import { CodeQLCliServer } from "../../../../../src/codeql-cli/cli";
+import {
+  KeyType,
+  resolveContextualQlPacksForDatabase,
+  resolveContextualQueries,
+} from "../../../../../src/language-support";
+import { CodeQLCliServer, DbInfo } from "../../../../../src/codeql-cli/cli";
 import { mockedObject } from "../../../utils/mocking.helpers";
+import * as queryResolver from "../../../../../src/local-queries/query-resolver";
+import { DatabaseItem } from "../../../../../src/databases/local-databases";
+import { Uri } from "vscode";
 
 describe("queryResolver", () => {
-  const resolveQueriesInSuite = jest.fn();
+  let qlpackOfDatabase: jest.SpiedFunction<
+    typeof queryResolver.qlpackOfDatabase
+  >;
+
+  const resolveQueriesInSuite: jest.MockedFunction<
+    typeof CodeQLCliServer.prototype.resolveQueriesInSuite
+  > = jest.fn();
+  const resolveDatabase: jest.MockedFunction<
+    typeof CodeQLCliServer.prototype.resolveDatabase
+  > = jest.fn();
+  const packDownload: jest.MockedFunction<
+    typeof CodeQLCliServer.prototype.packDownload
+  > = jest.fn();
 
   const mockCli = mockedObject<CodeQLCliServer>({
     resolveQueriesInSuite,
+    resolveDatabase,
+    packDownload,
   });
 
   beforeEach(() => {
+    qlpackOfDatabase = jest.spyOn(queryResolver, "qlpackOfDatabase");
+
     jest
       .spyOn(workspaceFolders, "getOnDiskWorkspaceFolders")
       .mockReturnValue([]);
     jest.spyOn(log, "showAndLogErrorMessage").mockResolvedValue(undefined);
   });
 
-  describe("resolveQueries", () => {
+  describe("resolveContextualQlPacksForDatabase", () => {
+    let databaseItem: DatabaseItem;
+
+    beforeEach(() => {
+      databaseItem = {
+        name: "my-db",
+        language: "csharp",
+        databaseUri: Uri.file("/a/b/c/db"),
+      } as DatabaseItem;
+    });
+
+    it("should resolve a qlpack when CLI returns qlpack", async () => {
+      qlpackOfDatabase.mockResolvedValue({
+        dbschemePack: "dbschemePack",
+        dbschemePackIsLibraryPack: false,
+      });
+
+      expect(
+        await resolveContextualQlPacksForDatabase(mockCli, databaseItem),
+      ).toEqual({
+        dbschemePack: "dbschemePack",
+        dbschemePackIsLibraryPack: false,
+      });
+    });
+
+    it("should return qlpack when downloading packs", async () => {
+      qlpackOfDatabase.mockRejectedValue(new Error("error"));
+      resolveDatabase.mockResolvedValue({
+        languages: ["csharp"],
+      } as DbInfo);
+
+      expect(
+        await resolveContextualQlPacksForDatabase(mockCli, databaseItem),
+      ).toEqual({
+        dbschemePack: "codeql/csharp-all",
+        dbschemePackIsLibraryPack: true,
+        queryPack: "codeql/csharp-queries",
+      });
+      expect(packDownload).toHaveBeenCalledTimes(1);
+      expect(packDownload).toHaveBeenCalledWith([
+        "codeql/csharp-all",
+        "codeql/csharp-queries",
+      ]);
+    });
+  });
+
+  describe("resolveContextualQueries", () => {
     it("should resolve a query", async () => {
-      resolveQueriesInSuite.mockReturnValue(["a", "b"]);
-      const result = await resolveQueries(
+      resolveQueriesInSuite.mockResolvedValue(["a", "b"]);
+      const result = await resolveContextualQueries(
         mockCli,
         { dbschemePack: "my-qlpack", dbschemePackIsLibraryPack: false },
         KeyType.DefinitionQuery,
@@ -53,10 +122,10 @@ describe("queryResolver", () => {
     });
 
     it("should throw an error when there are no queries found", async () => {
-      resolveQueriesInSuite.mockReturnValue([]);
+      resolveQueriesInSuite.mockResolvedValue([]);
 
       try {
-        await resolveQueries(
+        await resolveContextualQueries(
           mockCli,
           { dbschemePack: "my-qlpack", dbschemePackIsLibraryPack: false },
           KeyType.DefinitionQuery,