Unify model generation query running

2023-11-01 13:44:43 +01:00
parent 8a8a85fb9a
commit 9387d55263
4 changed files with 215 additions and 267 deletions
--- a/extensions/ql-vscode/src/model-editor/flow-model-queries.ts
+++ b/extensions/ql-vscode/src/model-editor/flow-model-queries.ts
@@ -1,21 +1,14 @@
-import { CancellationToken } from "vscode";
-import { DatabaseItem } from "../databases/local-databases";
 import { basename } from "path";
-import { QueryRunner } from "../query-server";
-import { CodeQLCliServer } from "../codeql-cli/cli";
+import { BaseLogger } from "../common/logging";
 import {
-  NotificationLogger,
-  showAndLogExceptionWithTelemetry,
-} from "../common/logging";
-import { getModelsAsDataLanguage } from "./languages";
-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 { telemetryListener } from "../common/vscode/telemetry";
-import { runQuery } from "../local-queries/run-query";
-import { resolveQueries } from "../local-queries";
+  getModelsAsDataLanguage,
+  ModelsAsDataLanguage,
+  ModelsAsDataLanguageModelType,
+} from "./languages";
+import { ModeledMethod } from "./modeled-method";
 import { QueryLanguage } from "../common/query-language";
+import { GenerateQueriesOptions, runGenerateQueries } from "./generate";
+import { DecodedBqrs } from "../common/bqrs-cli-types";

 const FLOW_MODEL_SUPPORTED_LANGUAGES = [
  QueryLanguage.CSharp,
@@ -28,158 +21,46 @@ export function isFlowModelGenerationSupported(
  return FLOW_MODEL_SUPPORTED_LANGUAGES.includes(language);
 }

-type FlowModelOptions = {
-  cliServer: CodeQLCliServer;
-  queryRunner: QueryRunner;
-  logger: NotificationLogger;
-  queryStorageDir: string;
-  databaseItem: DatabaseItem;
+type FlowModelOptions = GenerateQueriesOptions & {
+  logger: BaseLogger;
  language: QueryLanguage;
-  progress: ProgressCallback;
-  token: CancellationToken;
-  onResults: (results: ModeledMethod[]) => void | Promise<void>;
 };

-export async function runFlowModelQueries({
-  onResults,
-  ...options
-}: FlowModelOptions) {
-  const queries = await resolveFlowQueries(
-    options.cliServer,
-    options.databaseItem,
-  );
+const queriesToModel: Record<string, ModelsAsDataLanguageModelType> = {
+  "CaptureSummaryModels.ql": "summary",
+  "CaptureSinkModels.ql": "sink",
+  "CaptureSourceModels.ql": "source",
+  "CaptureNeutralModels.ql": "neutral",
+};

-  const queriesByBasename: Record<string, string> = {};
-  for (const query of queries) {
-    queriesByBasename[basename(query)] = query;
-  }
-
-  const summaryResults = await runSingleFlowQuery(
-    "summary",
-    queriesByBasename["CaptureSummaryModels.ql"],
-    0,
-    options,
-  );
-  if (summaryResults) {
-    await onResults(summaryResults);
-  }
-
-  const sinkResults = await runSingleFlowQuery(
-    "sink",
-    queriesByBasename["CaptureSinkModels.ql"],
-    1,
-    options,
-  );
-  if (sinkResults) {
-    await onResults(sinkResults);
-  }
-
-  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 resolveFlowQueries(
-  cliServer: CodeQLCliServer,
-  databaseItem: DatabaseItem,
-): Promise<string[]> {
-  const packsToSearch = [`codeql/${databaseItem.language}-queries`];
-
-  return await resolveQueries(
-    cliServer,
-    packsToSearch,
-    "flow model generator",
-    {
-      "tags contain": ["modelgenerator"],
-    },
-  );
-}
-
-async function runSingleFlowQuery(
-  type: Exclude<ModeledMethodType, "none">,
-  queryPath: string | undefined,
-  queryStep: number,
-  {
-    cliServer,
-    queryRunner,
-    logger,
-    queryStorageDir,
-    databaseItem,
-    language,
-    progress,
-    token,
-  }: Omit<FlowModelOptions, "onResults">,
-): Promise<ModeledMethod[]> {
-  // Check that the right query was found
-  if (queryPath === undefined) {
-    void showAndLogExceptionWithTelemetry(
-      logger,
-      telemetryListener,
-      redactableError`Failed to find ${type} query`,
-    );
-    return [];
-  }
-
-  // Run the query
-  const completedQuery = await runQuery({
-    queryRunner,
-    databaseItem,
-    queryPath,
-    queryStorageDir,
-    additionalPacks: getOnDiskWorkspaceFolders(),
-    extensionPacks: undefined,
-    progress: ({ step, message }) =>
-      progress({
-        message: `Generating ${type} model: ${message}`,
-        step: queryStep * 1000 + step,
-        maxStep: 4000,
-      }),
-    token,
-  });
-
-  if (!completedQuery) {
-    return [];
-  }
-
-  // Interpret the results
-  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
-
-  const definition = modelsAsDataLanguage.predicates[type];
-
-  const bqrsPath = completedQuery.outputDir.bqrsPath;
-
-  const bqrsInfo = await cliServer.bqrsInfo(bqrsPath);
-  if (bqrsInfo["result-sets"].length !== 1) {
-    void showAndLogExceptionWithTelemetry(
-      logger,
-      telemetryListener,
-      redactableError`Expected exactly one result set, got ${
-        bqrsInfo["result-sets"].length
-      } for ${basename(queryPath)}`,
+function parseFlowModelResults(
+  queryPath: string,
+  bqrs: DecodedBqrs,
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  logger: BaseLogger,
+): ModeledMethod[] {
+  if (Object.keys(bqrs).length !== 1) {
+    throw new Error(
+      `Expected exactly one result set from ${queryPath}, but got ${
+        Object.keys(bqrs).length
+      }`,
    );
  }

-  const resultSet = bqrsInfo["result-sets"][0];
+  const modelType = queriesToModel[basename(queryPath)];
+  if (!modelType) {
+    void logger.log(`Unknown model type for ${queryPath}`);
+    return [];
+  }

-  const decodedResults = await cliServer.bqrsDecode(bqrsPath, resultSet.name);
+  const resultSet = bqrs[Object.keys(bqrs)[0]];

-  const results = decodedResults.tuples;
+  const results = resultSet.tuples;
+
+  const definition = modelsAsDataLanguage.predicates[modelType];
+  if (!definition) {
+    throw new Error(`No definition for ${modelType}`);
+  }

  return (
    results
@@ -192,3 +73,37 @@ async function runSingleFlowQuery(
      })
  );
 }
+
+export async function runFlowModelQueries({
+  cliServer,
+  queryRunner,
+  logger,
+  queryStorageDir,
+  databaseItem,
+  language,
+  progress,
+  token,
+  onResults,
+}: FlowModelOptions) {
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
+  return runGenerateQueries(
+    {
+      queryConstraints: {
+        "tags contain": ["modelgenerator"],
+      },
+      filterQueries: (queryPath) => basename(queryPath) in queriesToModel,
+      parseResults: (queryPath, results) =>
+        parseFlowModelResults(queryPath, results, modelsAsDataLanguage, logger),
+    },
+    {
+      cliServer,
+      queryRunner,
+      queryStorageDir,
+      databaseItem,
+      progress,
+      token,
+      onResults,
+    },
+  );
+}
--- a/extensions/ql-vscode/src/model-editor/generate-model-queries.ts
+++ b/extensions/ql-vscode/src/model-editor/generate-model-queries.ts
@@ -1,21 +1,10 @@
-import { CancellationToken } from "vscode";
-import { DatabaseItem } from "../databases/local-databases";
-import { QueryRunner } from "../query-server";
-import { CodeQLCliServer } from "../codeql-cli/cli";
-import {
-  NotificationLogger,
-  showAndLogExceptionWithTelemetry,
-} from "../common/logging";
-import { getModelsAsDataLanguage } from "./languages";
-import { ProgressCallback } from "../common/vscode/progress";
-import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { BaseLogger, NotificationLogger } from "../common/logging";
+import { getModelsAsDataLanguage, ModelsAsDataLanguage } from "./languages";
 import { ModeledMethod } from "./modeled-method";
-import { redactableError } from "../common/errors";
-import { telemetryListener } from "../common/vscode/telemetry";
-import { runQuery } from "../local-queries/run-query";
-import { resolveQueries } from "../local-queries";
 import { QueryLanguage } from "../common/query-language";
 import { DataTuple } from "./model-extension-file";
+import { GenerateQueriesOptions, runGenerateQueries } from "./generate";
+import { DecodedBqrs } from "../common/bqrs-cli-types";

 const GENERATE_MODEL_SUPPORTED_LANGUAGES = [QueryLanguage.Ruby];

@@ -23,81 +12,19 @@ export function isGenerateModelSupported(language: QueryLanguage): boolean {
  return GENERATE_MODEL_SUPPORTED_LANGUAGES.includes(language);
 }

-type GenerateModelOptions = {
-  cliServer: CodeQLCliServer;
-  queryRunner: QueryRunner;
+type GenerateModelOptions = GenerateQueriesOptions & {
  logger: NotificationLogger;
-  queryStorageDir: string;
-  databaseItem: DatabaseItem;
  language: QueryLanguage;
-  progress: ProgressCallback;
-  token: CancellationToken;
 };

-// resolve (100) + query (1000) + interpret (100)
-const maxStep = 1200;
-
-export async function runGenerateModelQuery({
-  cliServer,
-  queryRunner,
-  logger,
-  queryStorageDir,
-  databaseItem,
-  language,
-  progress,
-  token,
-}: GenerateModelOptions): Promise<ModeledMethod[]> {
-  progress({
-    message: "Resolving generate model query",
-    step: 100,
-    maxStep,
-  });
-
-  const queryPath = await resolveGenerateModelQuery(
-    cliServer,
-    logger,
-    databaseItem,
-  );
-  if (queryPath === undefined) {
-    return [];
-  }
-
-  // Run the query
-  const completedQuery = await runQuery({
-    queryRunner,
-    databaseItem,
-    queryPath,
-    queryStorageDir,
-    additionalPacks: getOnDiskWorkspaceFolders(),
-    extensionPacks: undefined,
-    progress: ({ step, message }) =>
-      progress({
-        message: `Generating models: ${message}`,
-        step: 100 + step,
-        maxStep,
-      }),
-    token,
-  });
-
-  if (!completedQuery) {
-    return [];
-  }
-
-  progress({
-    message: "Decoding results",
-    step: 1100,
-    maxStep,
-  });
-
-  const decodedBqrs = await cliServer.bqrsDecodeAll(
-    completedQuery.outputDir.bqrsPath,
-  );
-
-  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
-
+function parseGenerateModelResults(
+  bqrs: DecodedBqrs,
+  modelsAsDataLanguage: ModelsAsDataLanguage,
+  logger: BaseLogger,
+): ModeledMethod[] {
  const modeledMethods: ModeledMethod[] = [];

-  for (const resultSetName in decodedBqrs) {
+  for (const resultSetName in bqrs) {
    const definition = Object.values(modelsAsDataLanguage.predicates).find(
      (definition) => definition.extensiblePredicate === resultSetName,
    );
@@ -107,7 +34,7 @@ export async function runGenerateModelQuery({
      continue;
    }

-    const resultSet = decodedBqrs[resultSetName];
+    const resultSet = bqrs[resultSetName];

    if (
      resultSet.tuples.some((tuple) =>
@@ -134,29 +61,21 @@ export async function runGenerateModelQuery({
  return modeledMethods;
 }

-async function resolveGenerateModelQuery(
-  cliServer: CodeQLCliServer,
-  logger: NotificationLogger,
-  databaseItem: DatabaseItem,
-): Promise<string | undefined> {
-  const packsToSearch = [`codeql/${databaseItem.language}-queries`];
+export async function runGenerateModelQuery({
+  logger,
+  language,
+  ...options
+}: GenerateModelOptions) {
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);

-  const queries = await resolveQueries(
-    cliServer,
-    packsToSearch,
-    "generate model",
+  return runGenerateQueries(
    {
-      "query path": "queries/modeling/GenerateModel.ql",
+      queryConstraints: {
+        "query path": "queries/modeling/GenerateModel.ql",
+      },
+      parseResults: (_queryPath, results) =>
+        parseGenerateModelResults(results, modelsAsDataLanguage, logger),
    },
+    options,
  );
-  if (queries.length !== 1) {
-    void showAndLogExceptionWithTelemetry(
-      logger,
-      telemetryListener,
-      redactableError`Expected exactly one generate model query, got ${queries.length}`,
-    );
-    return undefined;
-  }
-
-  return queries[0];
 }
--- a/extensions/ql-vscode/src/model-editor/generate.ts
+++ b/extensions/ql-vscode/src/model-editor/generate.ts
@@ -0,0 +1,113 @@
+import { CancellationToken } from "vscode";
+import { DatabaseItem } from "../databases/local-databases";
+import { basename } from "path";
+import { QueryRunner } from "../query-server";
+import { CodeQLCliServer } from "../codeql-cli/cli";
+import { ProgressCallback } from "../common/vscode/progress";
+import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
+import { ModeledMethod } from "./modeled-method";
+import { runQuery } from "../local-queries/run-query";
+import { QueryConstraints, resolveQueries } from "../local-queries";
+import { DecodedBqrs } from "../common/bqrs-cli-types";
+
+/**
+ * Options that are set by the caller of `runGenerateQueries`.
+ */
+type GenerateQueriesQueryOptions = {
+  queryConstraints: QueryConstraints;
+  filterQueries?: (queryPath: string) => boolean;
+  parseResults: (
+    queryPath: string,
+    results: DecodedBqrs,
+  ) => ModeledMethod[] | Promise<ModeledMethod[]>;
+};
+
+/**
+ * Options that are passed through by the caller of `runGenerateQueries`.
+ */
+export type GenerateQueriesOptions = {
+  cliServer: CodeQLCliServer;
+  queryRunner: QueryRunner;
+  queryStorageDir: string;
+  databaseItem: DatabaseItem;
+  progress: ProgressCallback;
+  token: CancellationToken;
+  onResults: (results: ModeledMethod[]) => void | Promise<void>;
+};
+
+export async function runGenerateQueries(
+  {
+    queryConstraints,
+    filterQueries,
+    parseResults,
+  }: GenerateQueriesQueryOptions,
+  { onResults, ...options }: GenerateQueriesOptions,
+) {
+  options.progress({
+    message: "Resolving queries",
+    step: 1,
+    maxStep: 5000,
+  });
+
+  const packsToSearch = [`codeql/${options.databaseItem.language}-queries`];
+  const queryPaths = await resolveQueries(
+    options.cliServer,
+    packsToSearch,
+    "generate model",
+    queryConstraints,
+  );
+
+  const filteredQueryPaths = filterQueries
+    ? queryPaths.filter(filterQueries)
+    : queryPaths;
+
+  const maxStep = filteredQueryPaths.length * 1000;
+
+  for (let i = 0; i < filteredQueryPaths.length; i++) {
+    const queryPath = filteredQueryPaths[i];
+
+    const bqrs = await runSingleGenerateQuery(queryPath, i, maxStep, options);
+    if (bqrs) {
+      await onResults(await parseResults(queryPath, bqrs));
+    }
+  }
+}
+
+async function runSingleGenerateQuery(
+  queryPath: string,
+  queryStep: number,
+  maxStep: number,
+  {
+    cliServer,
+    queryRunner,
+    queryStorageDir,
+    databaseItem,
+    progress,
+    token,
+  }: Omit<GenerateQueriesOptions, "onResults">,
+): Promise<DecodedBqrs | undefined> {
+  const queryBasename = basename(queryPath);
+
+  // Run the query
+  const completedQuery = await runQuery({
+    queryRunner,
+    databaseItem,
+    queryPath,
+    queryStorageDir,
+    additionalPacks: getOnDiskWorkspaceFolders(),
+    extensionPacks: undefined,
+    progress: ({ step, message }) =>
+      progress({
+        message: `Generating model from ${queryBasename}: ${message}`,
+        step: queryStep * 1000 + step,
+        maxStep,
+      }),
+    token,
+  });
+
+  if (!completedQuery) {
+    return undefined;
+  }
+
+  return cliServer.bqrsDecodeAll(completedQuery.outputDir.bqrsPath);
+}
--- a/extensions/ql-vscode/src/model-editor/model-editor-view.ts
+++ b/extensions/ql-vscode/src/model-editor/model-editor-view.ts
@@ -542,7 +542,7 @@ export class ModelEditorView extends AbstractWebview<
        const tokenSource = new CancellationTokenSource();

        try {
-          const modeledMethods = await runGenerateModelQuery({
+          await runGenerateModelQuery({
            cliServer: this.cliServer,
            queryRunner: this.queryRunner,
            logger: this.app.logger,
@@ -551,9 +551,10 @@ export class ModelEditorView extends AbstractWebview<
            language: this.language,
            progress,
            token: tokenSource.token,
+            onResults: async (modeledMethods) => {
+              this.addModeledMethodsFromArray(modeledMethods);
+            },
          });
-
-          this.addModeledMethodsFromArray(modeledMethods);
        } catch (e: unknown) {
          void showAndLogExceptionWithTelemetry(
            this.app.logger,