Resolve model editor queries from ql if present

Author: koesie10

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

@@ -39,12 +39,14 @@ export interface QueryConstraints {
  * @param cli The CLI instance to use.
  * @param qlpacks The list of packs to search.
  * @param constraints Constraints on the queries to search for.
+ * @param additionalPacks Additional pack paths to search.
  * @returns The found queries from the first pack in which any matching queries were found.
  */
 async function resolveQueriesFromPacks(
   cli: CodeQLCliServer,
   qlpacks: string[],
   constraints: QueryConstraints,
+  additionalPacks: string[] = [],
 ): Promise<string[]> {
   const suiteFile = (
     await file({
@@ -67,10 +69,10 @@ async function resolveQueriesFromPacks(
     "utf8",
   );
 
-  return await cli.resolveQueriesInSuite(
-    suiteFile,
-    getOnDiskWorkspaceFolders(),
-  );
+  return await cli.resolveQueriesInSuite(suiteFile, [
+    ...getOnDiskWorkspaceFolders(),
+    ...additionalPacks,
+  ]);
 }
 
 export async function resolveQueriesByLanguagePack(
@@ -97,18 +99,23 @@ export async function resolveQueriesByLanguagePack(
  * @param packsToSearch 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.
+ * @param allowNoQueriesFound If true, will not throw an error if no queries are found.
+ * @param additionalPacks Additional pack paths to search.
  * @returns The found queries from the first pack in which any matching queries were found.
  */
 export async function resolveQueries(
   cli: CodeQLCliServer,
   packsToSearch: string[],
   name: string,
   constraints: QueryConstraints,
+  allowNoQueriesFound = false,
+  additionalPacks: string[] = [],
 ): Promise<string[]> {
   const queries = await resolveQueriesFromPacks(
     cli,
     packsToSearch,
     constraints,
+    additionalPacks,
   );
   if (queries.length > 0) {
     return queries;
@@ -128,6 +135,10 @@ export async function resolveQueries(
     );
   }
 
+  if (allowNoQueriesFound) {
+    return [];
+  }
+
   const joinedPacksToSearch = packsToSearch.join(", ");
   const error = redactableError`No ${name} queries (${humanConstraints.join(
     ", ",

extensions/ql-vscode/src/model-editor/auto-model-codeml-queries.ts

@@ -7,7 +7,6 @@ import { Mode } from "./shared/mode";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { interpretResultsSarif } from "../query-results";
 import { join } from "path";
-import { assertNever } from "../common/helpers-pure";
 import { dir } from "tmp-promise";
 import { writeFile, outputFile } from "fs-extra";
 import { dump as dumpYaml } from "js-yaml";
@@ -16,17 +15,7 @@ import { runQuery } from "../local-queries/run-query";
 import { QueryMetadata } from "../common/interface-types";
 import { CancellationTokenSource } from "vscode";
 import { resolveQueries } from "../local-queries";
-
-function modeTag(mode: Mode): string {
-  switch (mode) {
-    case Mode.Application:
-      return "application-mode";
-    case Mode.Framework:
-      return "framework-mode";
-    default:
-      assertNever(mode);
-  }
-}
+import { modeTag } from "./mode-tag";
 
 type AutoModelQueriesOptions = {
   mode: Mode;

extensions/ql-vscode/src/model-editor/external-api-usage-queries.ts

@@ -16,6 +16,10 @@ import { fetchExternalApiQueries } from "./queries";
 import { Method } from "./method";
 import { runQuery } from "../local-queries/run-query";
 import { decodeBqrsToMethods } from "./bqrs";
+import {
+  resolveEndpointsQuery,
+  syntheticQueryPackName,
+} from "./model-editor-queries";
 
 type RunQueryOptions = {
   cliServer: CodeQLCliServer;
@@ -88,7 +92,24 @@ export async function runExternalApiQueries(
     await cliServer.resolveQlpacks(additionalPacks, true),
   );
 
-  const queryPath = join(queryDir, queryNameFromMode(mode));
+  progress({
+    message: "Resolving query",
+    step: 2,
+    maxStep: externalApiQueriesProgressMaxStep,
+  });
+
+  // Resolve the queries from either codeql/java-queries or from the temporary queryDir
+  const queryPath = await resolveEndpointsQuery(
+    cliServer,
+    databaseItem.language,
+    mode,
+    [syntheticQueryPackName],
+    [queryDir],
+    false,
+  );
+  if (!queryPath) {
+    return;
+  }
 
   // Run the actual query
   const completedQuery = await runQuery({

extensions/ql-vscode/src/model-editor/mode-tag.ts

@@ -0,0 +1,13 @@
+import { Mode } from "./shared/mode";
+import { assertNever } from "../common/helpers-pure";
+
+export function modeTag(mode: Mode): string {
+  switch (mode) {
+    case Mode.Application:
+      return "application-mode";
+    case Mode.Framework:
+      return "framework-mode";
+    default:
+      assertNever(mode);
+  }
+}

extensions/ql-vscode/src/model-editor/model-editor-module.ts

@@ -132,6 +132,7 @@ export class ModelEditorModule extends DisposableObject {
             const { path: queryDir, cleanup: cleanupQueryDir } = await dir({
               unsafeCleanup: true,
             });
+
             const success = await setUpPack(this.cliServer, queryDir, language);
             if (!success) {
               await cleanupQueryDir();

extensions/ql-vscode/src/model-editor/model-editor-queries.ts

@@ -5,9 +5,28 @@ import { dump } from "js-yaml";
 import { prepareExternalApiQuery } from "./external-api-usage-queries";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { showLlmGeneration } from "../config";
+import { Mode } from "./shared/mode";
+import { resolveQueries } from "../local-queries";
+import { modeTag } from "./mode-tag";
+import { extLogger } from "../common/logging/vscode";
+
+export const syntheticQueryPackName = "codeql/external-api-usage";
 
 /**
- * setUpPack sets up a directory to use for the data extension editor queries.
+ * setUpPack sets up a directory to use for the data extension editor queries if required.
+ *
+ * There are two cases (example language is Java):
+ * - In case the queries are present in the codeql/java-queries, we don't need to write our own queries
+ *   to disk. We still need to create a synthetic query pack so we can pass the queryDir to the query
+ *   resolver without caring about whether the queries are present in the pack or not.
+ * - In case the queries are not present in the codeql/java-queries, we need to write our own queries
+ *   to disk. We will create a synthetic query pack and install its dependencies so it is fully independent
+ *   and we can simply pass it through when resolving the queries.
+ *
+ * These steps together ensure that later steps of the process don't need to keep track of whether the queries
+ * are present in codeql/java-queries or in our own query pack. They just need to resolve the query.
+ *
+ * @param cliServer The CodeQL CLI server to use.
  * @param queryDir The directory to set up.
  * @param language The language to use for the queries.
  * @returns true if the setup was successful, false otherwise.
@@ -17,34 +36,111 @@ export async function setUpPack(
   queryDir: string,
   language: QueryLanguage,
 ): Promise<boolean> {
-  // Create the external API query
-  const externalApiQuerySuccess = await prepareExternalApiQuery(
-    queryDir,
+  // Download the required query packs
+  await cliServer.packDownload([`codeql/${language}-queries`]);
+
+  const applicationModeQuery = await resolveEndpointsQuery(
+    cliServer,
     language,
+    Mode.Application,
+    [],
+    [],
+    true,
   );
-  if (!externalApiQuerySuccess) {
-    return false;
-  }
 
-  // Set up a synthetic pack so that the query can be resolved later.
-  const syntheticQueryPack = {
-    name: "codeql/external-api-usage",
-    version: "0.0.0",
-    dependencies: {
-      [`codeql/${language}-all`]: "*",
-    },
-  };
+  if (applicationModeQuery) {
+    void extLogger.log("Using application mode queries");
 
-  const qlpackFile = join(queryDir, "codeql-pack.yml");
-  await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
-  await cliServer.packInstall(queryDir);
+    // Set up a synthetic pack so CodeQL doesn't crash later when we try
+    // to resolve a query within this directory
+    const syntheticQueryPack = {
+      name: syntheticQueryPackName,
+      version: "0.0.0",
+      dependencies: {},
+    };
 
-  // Install the other needed query packs
-  await cliServer.packDownload([`codeql/${language}-queries`]);
+    const qlpackFile = join(queryDir, "codeql-pack.yml");
+    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
+  } else {
+    void extLogger.log("Writing external API usage queries to disk");
 
+    // If we can't resolve the query, we need to write them to desk ourselves.
+    const externalApiQuerySuccess = await prepareExternalApiQuery(
+      queryDir,
+      language,
+    );
+    if (!externalApiQuerySuccess) {
+      return false;
+    }
+
+    // Set up a synthetic pack so that the query can be resolved later.
+    const syntheticQueryPack = {
+      name: syntheticQueryPackName,
+      version: "0.0.0",
+      dependencies: {
+        [`codeql/${language}-all`]: "*",
+      },
+    };
+
+    const qlpackFile = join(queryDir, "codeql-pack.yml");
+    await writeFile(qlpackFile, dump(syntheticQueryPack), "utf8");
+    await cliServer.packInstall(queryDir);
+  }
+
+  // Download any other required packs
   if (language === "java" && showLlmGeneration()) {
     await cliServer.packDownload([`codeql/${language}-automodel-queries`]);
   }
 
   return true;
 }
+
+/**
+ * Resolve the query path to the model editor endpoints query. All queries are tagged like this:
+ * modeleditor endpoints <mode>
+ * Example: modeleditor endpoints 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.
+ * @param allowNoQueriesFound If true, will not throw an error if no queries are found.
+ */
+export async function resolveEndpointsQuery(
+  cliServer: CodeQLCliServer,
+  language: string,
+  mode: Mode,
+  additionalPackNames: string[] = [],
+  additionalPackPaths: string[] = [],
+  allowNoQueriesFound = false,
+): Promise<string | undefined> {
+  const packsToSearch = [`codeql/${language}-queries`, ...additionalPackNames];
+
+  // First, resolve the query that we want to run.
+  // All queries are tagged like this:
+  // internal extract automodel <mode> <queryTag>
+  // Example: internal extract automodel framework-mode candidates
+  const queries = await resolveQueries(
+    cliServer,
+    packsToSearch,
+    `Fetch endpoints query for ${mode}`,
+    {
+      kind: "table",
+      "tags contain all": ["modeleditor", "endpoints", modeTag(mode)],
+    },
+    allowNoQueriesFound,
+    additionalPackPaths,
+  );
+  if (queries.length > 1) {
+    throw new Error(
+      `Found multiple endpoints queries for ${mode}. Can't continue`,
+    );
+  }
+
+  if (queries.length === 0) {
+    return undefined;
+  }
+
+  return queries[0];
+}

extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/external-api-usage-query.test.ts

@@ -46,6 +46,9 @@ describe("external api usage query", () => {
           resolveQlpacks: jest.fn().mockResolvedValue({
             "my/extensions": "/a/b/c/",
           }),
+          resolveQueriesInSuite: jest
+            .fn()
+            .mockResolvedValue(["/a/b/c/ApplicationModeEndpoints.ql"]),
           packPacklist: jest
             .fn()
             .mockResolvedValue([
@@ -104,6 +107,9 @@ describe("external api usage query", () => {
           resolveQlpacks: jest.fn().mockResolvedValue({
             "my/extensions": "/a/b/c/",
           }),
+          resolveQueriesInSuite: jest
+            .fn()
+            .mockResolvedValue(["/a/b/c/ApplicationModeEndpoints.ql"]),
           packPacklist: jest
             .fn()
             .mockResolvedValue([