diff --git a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts b/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts index 2fcb5c0f015..e1b14e8465b 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts +++ b/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, diff --git a/extensions/ql-vscode/src/data-extensions-editor/flow-model-queries.ts b/extensions/ql-vscode/src/data-extensions-editor/flow-model-queries.ts index 0adbe7b7bbb..a9a538f9260 100644 --- a/extensions/ql-vscode/src/data-extensions-editor/flow-model-queries.ts +++ b/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; }; -async function resolveQueries( - cliServer: CodeQLCliServer, - databaseItem: DatabaseItem, -): Promise { - 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 = {}; + 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 { + const qlpacks = await qlpackOfDatabase(cliServer, databaseItem); + + return await resolveQueries(cliServer, qlpacks, "flow model generator", { + "tags contain": ["modelgenerator"], + }); +} + +async function runSingleFlowQuery( type: Exclude, queryPath: string | undefined, queryStep: number, @@ -77,6 +103,7 @@ async function getModeledMethodsFromFlow( token, }: Omit, ): Promise { + // 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 = {}; - 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); - } -}