Unify model generation query running
This commit is contained in:
Koen Vlaswinkel
@@ -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,159 +21,47 @@ 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;
|
||||
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 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`,
|
||||
);
|
||||
const modelType = queriesToModel[basename(queryPath)];
|
||||
if (!modelType) {
|
||||
void logger.log(`Unknown model type for ${queryPath}`);
|
||||
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,
|
||||
});
|
||||
const resultSet = bqrs[Object.keys(bqrs)[0]];
|
||||
|
||||
if (!completedQuery) {
|
||||
return [];
|
||||
const results = resultSet.tuples;
|
||||
|
||||
const definition = modelsAsDataLanguage.predicates[modelType];
|
||||
if (!definition) {
|
||||
throw new Error(`No definition for ${modelType}`);
|
||||
}
|
||||
|
||||
// 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)}`,
|
||||
);
|
||||
}
|
||||
|
||||
const resultSet = bqrsInfo["result-sets"][0];
|
||||
|
||||
const decodedResults = await cliServer.bqrsDecode(bqrsPath, resultSet.name);
|
||||
|
||||
const results = decodedResults.tuples;
|
||||
|
||||
return (
|
||||
results
|
||||
// This is just a sanity check. The query should only return strings.
|
||||
@@ -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,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
{
|
||||
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];
|
||||
}
|
||||
|
||||
113
extensions/ql-vscode/src/model-editor/generate.ts
Normal file
113
extensions/ql-vscode/src/model-editor/generate.ts
Normal file
@@ -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);
|
||||
}
|
||||
@@ -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);
|
||||
},
|
||||
});
|
||||
} catch (e: unknown) {
|
||||
void showAndLogExceptionWithTelemetry(
|
||||
this.app.logger,
|
||||
|
||||
Reference in New Issue
Block a user