Merge pull request #2742 from github/starcke/flow-queries-refactor

Starcke/flow queries refactor

Author: starcke

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

@@ -22,7 +22,7 @@ import {
 import { DatabaseItem, DatabaseManager } from "../databases/local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { asError, assertNever, getErrorMessage } from "../common/helpers-pure";
-import { generateFlowModel } from "./flow-model-queries";
+import { runFlowModelQueries } from "./flow-model-queries";
 import { promptImportGithubDatabase } from "../databases/database-fetcher";
 import { App } from "../common/app";
 import { showResolvableLocation } from "../databases/local-databases/locations";
@@ -389,7 +389,7 @@ export class DataExtensionsEditorView extends AbstractWebview<
         });
 
         try {
-          await generateFlowModel({
+          await runFlowModelQueries({
             cliServer: this.cliServer,
             queryRunner: this.queryRunner,
             queryStorageDir: this.queryStorageDir,

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

@@ -3,19 +3,16 @@ import { DatabaseItem } from "../databases/local-databases";
 import { basename } from "path";
 import { QueryRunner } from "../query-server";
 import { CodeQLCliServer } from "../codeql-cli/cli";
-import { showAndLogExceptionWithTelemetry, TeeLogger } from "../common/logging";
+import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
 import { extensiblePredicateDefinitions } from "./predicates";
 import { ProgressCallback } from "../common/vscode/progress";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { ModeledMethod, ModeledMethodType } from "./modeled-method";
 import { redactableError } from "../common/errors";
-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 "../local-queries";
+import { qlpackOfDatabase, resolveQueries } from "../local-queries";
 import { telemetryListener } from "../common/vscode/telemetry";
+import { runQuery } from "../local-queries/run-query";
 
 type FlowModelOptions = {
   cliServer: CodeQLCliServer;
@@ -27,44 +24,73 @@ type FlowModelOptions = {
   onResults: (results: ModeledMethod[]) => void | Promise<void>;
 };
 
-async function resolveQueries(
-  cliServer: CodeQLCliServer,
-  databaseItem: DatabaseItem,
-): Promise<string[]> {
-  const qlpacks = await qlpackOfDatabase(cliServer, databaseItem);
+export async function runFlowModelQueries({
+  onResults,
+  ...options
+}: FlowModelOptions) {
+  const queries = await resolveFlowQueries(
+    options.cliServer,
+    options.databaseItem,
+  );
 
-  const packsToSearch: string[] = [];
+  const queriesByBasename: Record<string, string> = {};
+  for (const query of queries) {
+    queriesByBasename[basename(query)] = query;
+  }
 
-  // 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 summaryResults = await runSingleFlowQuery(
+    "summary",
+    queriesByBasename["CaptureSummaryModels.ql"],
+    0,
+    options,
+  );
+  if (summaryResults) {
+    await onResults(summaryResults);
   }
 
-  const suiteFile = (
-    await file({
-      postfix: ".qls",
-    })
-  ).path;
-  const suiteYaml = [];
-  for (const qlpack of packsToSearch) {
-    suiteYaml.push({
-      from: qlpack,
-      queries: ".",
-      include: {
-        "tags contain": "modelgenerator",
-      },
-    });
+  const sinkResults = await runSingleFlowQuery(
+    "sink",
+    queriesByBasename["CaptureSinkModels.ql"],
+    1,
+    options,
+  );
+  if (sinkResults) {
+    await onResults(sinkResults);
   }
-  await writeFile(suiteFile, dump(suiteYaml), "utf8");
 
-  return await cliServer.resolveQueriesInSuite(
-    suiteFile,
-    getOnDiskWorkspaceFolders(),
+  const sourceResults = await runSingleFlowQuery(
+    "source",
+    queriesByBasename["CaptureSourceModels.ql"],
+    2,
+    options,
+  );
+  if (sourceResults) {
+    await onResults(sourceResults);
+  }
+
+  const neutralResults = await runSingleFlowQuery(
+    "neutral",
+    queriesByBasename["CaptureNeutralModels.ql"],
+    3,
+    options,
   );
+  if (neutralResults) {
+    await onResults(neutralResults);
+  }
 }
 
-async function getModeledMethodsFromFlow(
+async function resolveFlowQueries(
+  cliServer: CodeQLCliServer,
+  databaseItem: DatabaseItem,
+): Promise<string[]> {
+  const qlpacks = await qlpackOfDatabase(cliServer, databaseItem);
+
+  return await resolveQueries(cliServer, qlpacks, "flow model generator", {
+    "tags contain": ["modelgenerator"],
+  });
+}
+
+async function runSingleFlowQuery(
   type: Exclude<ModeledMethodType, "none">,
   queryPath: string | undefined,
   queryStep: number,
@@ -77,6 +103,7 @@ async function getModeledMethodsFromFlow(
     token,
   }: Omit<FlowModelOptions, "onResults">,
 ): Promise<ModeledMethod[]> {
+  // Check that the right query was found
   if (queryPath === undefined) {
     void showAndLogExceptionWithTelemetry(
       extLogger,
@@ -86,45 +113,32 @@ async function getModeledMethodsFromFlow(
     return [];
   }
 
-  const definition = extensiblePredicateDefinitions[type];
-
-  const queryRun = queryRunner.createQueryRun(
-    databaseItem.databaseUri.fsPath,
-    {
-      queryPath,
-      quickEvalPosition: undefined,
-      quickEvalCountOnly: false,
-    },
-    false,
-    getOnDiskWorkspaceFolders(),
-    undefined,
+  // Run the query
+  const completedQuery = await runQuery({
+    cliServer,
+    queryRunner,
+    databaseItem,
+    queryPath,
     queryStorageDir,
-    undefined,
-    undefined,
-  );
-
-  const queryResult = await queryRun.evaluate(
-    ({ step, message }) =>
+    additionalPacks: getOnDiskWorkspaceFolders(),
+    extensionPacks: undefined,
+    progress: ({ step, message }) =>
       progress({
         message: `Generating ${type} model: ${message}`,
         step: queryStep * 1000 + step,
         maxStep: 4000,
       }),
     token,
-    new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
-  );
-  if (queryResult.resultType !== QueryResultType.SUCCESS) {
-    void showAndLogExceptionWithTelemetry(
-      extLogger,
-      telemetryListener,
-      redactableError`Failed to run ${basename(queryPath)} query: ${
-        queryResult.message ?? "No message"
-      }`,
-    );
+  });
+
+  if (!completedQuery) {
     return [];
   }
 
-  const bqrsPath = queryResult.outputDir.bqrsPath;
+  // Interpret the results
+  const definition = extensiblePredicateDefinitions[type];
+
+  const bqrsPath = completedQuery.outputDir.bqrsPath;
 
   const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
   if (bqrsInfo["result-sets"].length !== 1) {
@@ -154,55 +168,3 @@ async function getModeledMethodsFromFlow(
       })
   );
 }
-
-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",
-    queriesByBasename["CaptureSummaryModels.ql"],
-    0,
-    options,
-  );
-  if (summaryResults) {
-    await onResults(summaryResults);
-  }
-
-  const sinkResults = await getModeledMethodsFromFlow(
-    "sink",
-    queriesByBasename["CaptureSinkModels.ql"],
-    1,
-    options,
-  );
-  if (sinkResults) {
-    await onResults(sinkResults);
-  }
-
-  const sourceResults = await getModeledMethodsFromFlow(
-    "source",
-    queriesByBasename["CaptureSourceModels.ql"],
-    2,
-    options,
-  );
-  if (sourceResults) {
-    await onResults(sourceResults);
-  }
-
-  const neutralResults = await getModeledMethodsFromFlow(
-    "neutral",
-    queriesByBasename["CaptureNeutralModels.ql"],
-    3,
-    options,
-  );
-  if (neutralResults) {
-    await onResults(neutralResults);
-  }
-}