Merge pull request #2416 from github/koesie10/resolve-data-extensions-editor-queries

Resolve queries instead of hardcoding the path

Author: koesie10

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -5,7 +5,6 @@ import {
   ViewColumn,
   window,
   workspace,
-  WorkspaceFolder,
 } from "vscode";
 import {
   AbstractWebview,
@@ -47,19 +46,6 @@ import {
 import { showLlmGeneration } from "../config";
 import { getAutoModelUsages } from "./auto-model-usages-query";
 
-function getQlSubmoduleFolder(): WorkspaceFolder | undefined {
-  const workspaceFolder = workspace.workspaceFolders?.find(
-    (folder) => folder.name === "ql",
-  );
-  if (!workspaceFolder) {
-    void extLogger.log("No workspace folder 'ql' found");
-
-    return;
-  }
-
-  return workspaceFolder;
-}
-
 export class DataExtensionsEditorView extends AbstractWebview<
   ToDataExtensionsEditorMessage,
   FromDataExtensionsEditorMessage
@@ -324,11 +310,6 @@ export class DataExtensionsEditorView extends AbstractWebview<
     // but we need to set it back to the originally selected database.
     await this.databaseManager.setCurrentDatabaseItem(selectedDatabase);
 
-    const workspaceFolder = getQlSubmoduleFolder();
-    if (!workspaceFolder) {
-      return;
-    }
-
     await this.showProgress({
       step: 0,
       maxStep: 4000,
@@ -340,7 +321,6 @@ export class DataExtensionsEditorView extends AbstractWebview<
         cliServer: this.cliServer,
         queryRunner: this.queryRunner,
         queryStorageDir: this.queryStorageDir,
-        qlDir: workspaceFolder.uri.fsPath,
         databaseItem: database,
         onResults: async (results) => {
           const modeledMethodsByName: Record<string, ModeledMethod> = {};

extensions/ql-vscode/src/data-extensions-editor/generate-flow-model.ts

@@ -1,6 +1,6 @@
 import { CancellationToken } from "vscode";
 import { DatabaseItem } from "../databases/local-databases";
-import { join } from "path";
+import { basename } from "path";
 import { QueryRunner } from "../query-server";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { TeeLogger } from "../common";
@@ -16,45 +16,84 @@ import {
 } from "./modeled-method";
 import { redactableError } from "../pure/errors";
 import { QueryResultType } from "../pure/new-messages";
+import { file } from "tmp-promise";
+import { writeFile } from "fs-extra";
+import { dump } from "js-yaml";
+import { qlpackOfDatabase } from "../language-support";
 
 type FlowModelOptions = {
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
   queryStorageDir: string;
-  qlDir: string;
   databaseItem: DatabaseItem;
   progress: ProgressCallback;
   token: CancellationToken;
   onResults: (results: ModeledMethodWithSignature[]) => void | Promise<void>;
 };
 
+async function resolveQueries(
+  cliServer: CodeQLCliServer,
+  databaseItem: DatabaseItem,
+): Promise<string[]> {
+  const qlpacks = await qlpackOfDatabase(cliServer, databaseItem);
+
+  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 suiteFile = (
+    await file({
+      postfix: ".qls",
+    })
+  ).path;
+  const suiteYaml = [];
+  for (const qlpack of packsToSearch) {
+    suiteYaml.push({
+      from: qlpack,
+      queries: ".",
+      include: {
+        "tags contain": "modelgenerator",
+      },
+    });
+  }
+  await writeFile(suiteFile, dump(suiteYaml), "utf8");
+
+  return await cliServer.resolveQueriesInSuite(
+    suiteFile,
+    getOnDiskWorkspaceFolders(),
+  );
+}
+
 async function getModeledMethodsFromFlow(
   type: Exclude<ModeledMethodType, "none">,
-  queryName: string,
+  queryPath: string | undefined,
   queryStep: number,
   {
     cliServer,
     queryRunner,
     queryStorageDir,
-    qlDir,
     databaseItem,
     progress,
     token,
   }: Omit<FlowModelOptions, "onResults">,
 ): Promise<ModeledMethodWithSignature[]> {
-  const definition = extensiblePredicateDefinitions[type];
+  if (queryPath === undefined) {
+    void showAndLogExceptionWithTelemetry(
+      redactableError`Failed to find ${type} query`,
+    );
+    return [];
+  }
 
-  const query = join(
-    qlDir,
-    databaseItem.language,
-    "ql/src/utils/modelgenerator",
-    queryName,
-  );
+  const definition = extensiblePredicateDefinitions[type];
 
   const queryRun = queryRunner.createQueryRun(
     databaseItem.databaseUri.fsPath,
     {
-      queryPath: query,
+      queryPath,
       quickEvalPosition: undefined,
       quickEvalCountOnly: false,
     },
@@ -78,7 +117,7 @@ async function getModeledMethodsFromFlow(
   );
   if (queryResult.resultType !== QueryResultType.SUCCESS) {
     void showAndLogExceptionWithTelemetry(
-      redactableError`Failed to run ${queryName} query: ${
+      redactableError`Failed to run ${basename(queryPath)} query: ${
         queryResult.message ?? "No message"
       }`,
     );
@@ -90,7 +129,9 @@ async function getModeledMethodsFromFlow(
   const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
   if (bqrsInfo["result-sets"].length !== 1) {
     void showAndLogExceptionWithTelemetry(
-      redactableError`Expected exactly one result set, got ${bqrsInfo["result-sets"].length} for ${queryName}`,
+      redactableError`Expected exactly one result set, got ${
+        bqrsInfo["result-sets"].length
+      } for ${basename(queryPath)}`,
     );
   }
 
@@ -116,9 +157,16 @@ export async function generateFlowModel({
   onResults,
   ...options
 }: FlowModelOptions) {
+  const queries = await resolveQueries(options.cliServer, options.databaseItem);
+
+  const queriesByBasename: Record<string, string> = {};
+  for (const query of queries) {
+    queriesByBasename[basename(query)] = query;
+  }
+
   const summaryResults = await getModeledMethodsFromFlow(
     "summary",
-    "CaptureSummaryModels.ql",
+    queriesByBasename["CaptureSummaryModels.ql"],
     0,
     options,
   );
@@ -128,7 +176,7 @@ export async function generateFlowModel({
 
   const sinkResults = await getModeledMethodsFromFlow(
     "sink",
-    "CaptureSinkModels.ql",
+    queriesByBasename["CaptureSinkModels.ql"],
     1,
     options,
   );
@@ -138,7 +186,7 @@ export async function generateFlowModel({
 
   const sourceResults = await getModeledMethodsFromFlow(
     "source",
-    "CaptureSourceModels.ql",
+    queriesByBasename["CaptureSourceModels.ql"],
     2,
     options,
   );
@@ -148,7 +196,7 @@ export async function generateFlowModel({
 
   const neutralResults = await getModeledMethodsFromFlow(
     "neutral",
-    "CaptureNeutralModels.ql",
+    queriesByBasename["CaptureNeutralModels.ql"],
     3,
     options,
   );