Merge pull request #3057 from github/koesie10/contextual-no-submodule

koesie10 · web-flow · commit bc2847a12e0a · 2023-11-13T09:41:06.000+01:00
Make contextual queries work for fresh installs
diff --git a/extensions/ql-vscode/src/language-support/contextual/location-finder.ts b/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,
diff --git a/extensions/ql-vscode/src/language-support/contextual/query-resolver.ts b/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(
diff --git a/extensions/ql-vscode/src/language-support/contextual/template-provider.ts b/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,
diff --git a/extensions/ql-vscode/src/local-queries/query-resolver.ts b/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">,
diff --git a/extensions/ql-vscode/test/vscode-tests/cli-integration/run-cli.test.ts b/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.
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,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,
