Skip to content

Commit 8a8a85f

Browse files
koesie10koesie10
authored
Merge pull request #3033 from github/koesie10/generate-model
Add generation of Ruby models
2 parents 978d8d3 + 2dbc50e commit 8a8a85f

File tree

14 files changed

+433
-30
lines changed

14 files changed

+433
-30
lines changed

‎extensions/ql-vscode/src/codeql-cli/cli.ts‎

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,11 @@ import tk from "tree-kill";
1010
import { promisify } from "util";
1111
import { CancellationToken, Disposable, Uri } from "vscode";
1212

13-
import { BQRSInfo, DecodedBqrsChunk } from "../common/bqrs-cli-types";
13+
import {
14+
BQRSInfo,
15+
DecodedBqrs,
16+
DecodedBqrsChunk,
17+
} from "../common/bqrs-cli-types";
1418
import { allowCanaryQueryServer, CliConfig } from "../config";
1519
import {
1620
DistributionProvider,
@@ -1040,6 +1044,18 @@ export class CodeQLCliServer implements Disposable {
10401044
);
10411045
}
10421046

1047+
/**
1048+
* Gets all results from a bqrs.
1049+
* @param bqrsPath The path to the bqrs.
1050+
*/
1051+
async bqrsDecodeAll(bqrsPath: string): Promise<DecodedBqrs> {
1052+
return await this.runJsonCodeQlCliCommand<DecodedBqrs>(
1053+
["bqrs", "decode"],
1054+
[bqrsPath],
1055+
"Reading all bqrs data",
1056+
);
1057+
}
1058+
10431059
async runInterpretCommand(
10441060
format: string,
10451061
additonalArgs: string[],

‎extensions/ql-vscode/src/common/bqrs-cli-types.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,3 +121,5 @@ export interface DecodedBqrsChunk {
121121
next?: number;
122122
columns: BqrsColumn[];
123123
}
124+
125+
export type DecodedBqrs = Record<string, DecodedBqrsChunk>;

‎extensions/ql-vscode/src/local-queries/query-resolver.ts‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@ export interface QueryConstraints {
3131
kind?: string;
3232
"tags contain"?: string[];
3333
"tags contain all"?: string[];
34+
"query filename"?: string;
35+
"query path"?: string;
3436
}
3537

3638
/**
@@ -132,6 +134,14 @@ export async function resolveQueries(
132134
`tagged all of "${constraints["tags contain all"].join(" ")}"`,
133135
);
134136
}
137+
if (constraints["query filename"] !== undefined) {
138+
humanConstraints.push(
139+
`with query filename "${constraints["query filename"]}"`,
140+
);
141+
}
142+
if (constraints["query path"] !== undefined) {
143+
humanConstraints.push(`with query path "${constraints["query path"]}"`);
144+
}
135145

136146
const joinedPacksToSearch = packsToSearch.join(", ");
137147
const error = redactableError`No ${name} queries (${humanConstraints.join(

‎extensions/ql-vscode/src/model-editor/generate-model-queries.ts‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
import { CancellationToken } from "vscode";
2+
import { DatabaseItem } from "../databases/local-databases";
3+
import { QueryRunner } from "../query-server";
4+
import { CodeQLCliServer } from "../codeql-cli/cli";
5+
import {
6+
NotificationLogger,
7+
showAndLogExceptionWithTelemetry,
8+
} from "../common/logging";
9+
import { getModelsAsDataLanguage } from "./languages";
10+
import { ProgressCallback } from "../common/vscode/progress";
11+
import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
12+
import { ModeledMethod } from "./modeled-method";
13+
import { redactableError } from "../common/errors";
14+
import { telemetryListener } from "../common/vscode/telemetry";
15+
import { runQuery } from "../local-queries/run-query";
16+
import { resolveQueries } from "../local-queries";
17+
import { QueryLanguage } from "../common/query-language";
18+
import { DataTuple } from "./model-extension-file";
19+
20+
const GENERATE_MODEL_SUPPORTED_LANGUAGES = [QueryLanguage.Ruby];
21+
22+
export function isGenerateModelSupported(language: QueryLanguage): boolean {
23+
return GENERATE_MODEL_SUPPORTED_LANGUAGES.includes(language);
24+
}
25+
26+
type GenerateModelOptions = {
27+
cliServer: CodeQLCliServer;
28+
queryRunner: QueryRunner;
29+
logger: NotificationLogger;
30+
queryStorageDir: string;
31+
databaseItem: DatabaseItem;
32+
language: QueryLanguage;
33+
progress: ProgressCallback;
34+
token: CancellationToken;
35+
};
36+
37+
// resolve (100) + query (1000) + interpret (100)
38+
const maxStep = 1200;
39+
40+
export async function runGenerateModelQuery({
41+
cliServer,
42+
queryRunner,
43+
logger,
44+
queryStorageDir,
45+
databaseItem,
46+
language,
47+
progress,
48+
token,
49+
}: GenerateModelOptions): Promise<ModeledMethod[]> {
50+
progress({
51+
message: "Resolving generate model query",
52+
step: 100,
53+
maxStep,
54+
});
55+
56+
const queryPath = await resolveGenerateModelQuery(
57+
cliServer,
58+
logger,
59+
databaseItem,
60+
);
61+
if (queryPath === undefined) {
62+
return [];
63+
}
64+
65+
// Run the query
66+
const completedQuery = await runQuery({
67+
queryRunner,
68+
databaseItem,
69+
queryPath,
70+
queryStorageDir,
71+
additionalPacks: getOnDiskWorkspaceFolders(),
72+
extensionPacks: undefined,
73+
progress: ({ step, message }) =>
74+
progress({
75+
message: `Generating models: ${message}`,
76+
step: 100 + step,
77+
maxStep,
78+
}),
79+
token,
80+
});
81+
82+
if (!completedQuery) {
83+
return [];
84+
}
85+
86+
progress({
87+
message: "Decoding results",
88+
step: 1100,
89+
maxStep,
90+
});
91+
92+
const decodedBqrs = await cliServer.bqrsDecodeAll(
93+
completedQuery.outputDir.bqrsPath,
94+
);
95+
96+
const modelsAsDataLanguage = getModelsAsDataLanguage(language);
97+
98+
const modeledMethods: ModeledMethod[] = [];
99+
100+
for (const resultSetName in decodedBqrs) {
101+
const definition = Object.values(modelsAsDataLanguage.predicates).find(
102+
(definition) => definition.extensiblePredicate === resultSetName,
103+
);
104+
if (definition === undefined) {
105+
void logger.log(`No predicate found for ${resultSetName}`);
106+
107+
continue;
108+
}
109+
110+
const resultSet = decodedBqrs[resultSetName];
111+
112+
if (
113+
resultSet.tuples.some((tuple) =>
114+
tuple.some((value) => typeof value === "object"),
115+
)
116+
) {
117+
void logger.log(
118+
`Skipping ${resultSetName} because it contains undefined values`,
119+
);
120+
continue;
121+
}
122+
123+
modeledMethods.push(
124+
...resultSet.tuples.map((tuple) => {
125+
const row = tuple.filter(
126+
(value): value is DataTuple => typeof value !== "object",
127+
);
128+
129+
return definition.readModeledMethod(row);
130+
}),
131+
);
132+
}
133+
134+
return modeledMethods;
135+
}
136+
137+
async function resolveGenerateModelQuery(
138+
cliServer: CodeQLCliServer,
139+
logger: NotificationLogger,
140+
databaseItem: DatabaseItem,
141+
): Promise<string | undefined> {
142+
const packsToSearch = [`codeql/${databaseItem.language}-queries`];
143+
144+
const queries = await resolveQueries(
145+
cliServer,
146+
packsToSearch,
147+
"generate model",
148+
{
149+
"query path": "queries/modeling/GenerateModel.ql",
150+
},
151+
);
152+
if (queries.length !== 1) {
153+
void showAndLogExceptionWithTelemetry(
154+
logger,
155+
telemetryListener,
156+
redactableError`Expected exactly one generate model query, got ${queries.length}`,
157+
);
158+
return undefined;
159+
}
160+
161+
return queries[0];
162+
}

‎extensions/ql-vscode/src/model-editor/model-editor-view.ts‎

Lines changed: 61 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ import { ModelingStore } from "./modeling-store";
5151
import { ModelEditorViewTracker } from "./model-editor-view-tracker";
5252
import { ModelingEvents } from "./modeling-events";
5353
import { getModelsAsDataLanguage, ModelsAsDataLanguage } from "./languages";
54+
import {
55+
isGenerateModelSupported,
56+
runGenerateModelQuery,
57+
} from "./generate-model-queries";
5458

5559
export class ModelEditorView extends AbstractWebview<
5660
ToModelEditorMessage,
@@ -266,7 +270,11 @@ export class ModelEditorView extends AbstractWebview<
266270

267271
break;
268272
case "generateMethod":
269-
await this.generateModeledMethods();
273+
if (isFlowModelGenerationSupported(this.language)) {
274+
await this.generateModeledMethodsFromFlow();
275+
} else if (isGenerateModelSupported(this.language)) {
276+
await this.generateModeledMethodsFromGenerateModel();
277+
}
270278
void telemetryListener?.sendUIInteraction(
271279
"model-editor-generate-modeled-methods",
272280
);
@@ -369,9 +377,10 @@ export class ModelEditorView extends AbstractWebview<
369377
}
370378

371379
private async setViewState(): Promise<void> {
372-
const showFlowGeneration =
380+
const showGenerateButton =
373381
this.modelConfig.flowGeneration &&
374-
isFlowModelGenerationSupported(this.language);
382+
(isFlowModelGenerationSupported(this.language) ||
383+
isGenerateModelSupported(this.language));
375384

376385
const showLlmButton =
377386
this.databaseItem.language === "java" && this.modelConfig.llmGeneration;
@@ -388,7 +397,7 @@ export class ModelEditorView extends AbstractWebview<
388397
viewState: {
389398
extensionPack: this.extensionPack,
390399
language: this.language,
391-
showFlowGeneration,
400+
showGenerateButton,
392401
showLlmButton,
393402
showMultipleModels: this.modelConfig.showMultipleModels,
394403
mode: this.modelingStore.getMode(this.databaseItem),
@@ -465,7 +474,7 @@ export class ModelEditorView extends AbstractWebview<
465474
}
466475
}
467476

468-
protected async generateModeledMethods(): Promise<void> {
477+
protected async generateModeledMethodsFromFlow(): Promise<void> {
469478
await withProgress(
470479
async (progress) => {
471480
const tokenSource = new CancellationTokenSource();
@@ -508,19 +517,7 @@ export class ModelEditorView extends AbstractWebview<
508517
databaseItem: addedDatabase ?? this.databaseItem,
509518
language: this.language,
510519
onResults: async (modeledMethods) => {
511-
const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
512-
513-
for (const modeledMethod of modeledMethods) {
514-
if (!(modeledMethod.signature in modeledMethodsByName)) {
515-
modeledMethodsByName[modeledMethod.signature] = [];
516-
}
517-
518-
modeledMethodsByName[modeledMethod.signature].push(
519-
modeledMethod,
520-
);
521-
}
522-
523-
this.addModeledMethods(modeledMethodsByName);
520+
this.addModeledMethodsFromArray(modeledMethods);
524521
},
525522
progress,
526523
token: tokenSource.token,
@@ -539,6 +536,38 @@ export class ModelEditorView extends AbstractWebview<
539536
);
540537
}
541538

539+
protected async generateModeledMethodsFromGenerateModel(): Promise<void> {
540+
await withProgress(
541+
async (progress) => {
542+
const tokenSource = new CancellationTokenSource();
543+
544+
try {
545+
const modeledMethods = await runGenerateModelQuery({
546+
cliServer: this.cliServer,
547+
queryRunner: this.queryRunner,
548+
logger: this.app.logger,
549+
queryStorageDir: this.queryStorageDir,
550+
databaseItem: this.databaseItem,
551+
language: this.language,
552+
progress,
553+
token: tokenSource.token,
554+
});
555+
556+
this.addModeledMethodsFromArray(modeledMethods);
557+
} catch (e: unknown) {
558+
void showAndLogExceptionWithTelemetry(
559+
this.app.logger,
560+
this.app.telemetry,
561+
redactableError(
562+
asError(e),
563+
)`Failed to generate models: ${getErrorMessage(e)}`,
564+
);
565+
}
566+
},
567+
{ cancellable: false },
568+
);
569+
}
570+
542571
private async generateModeledMethodsFromLlm(
543572
packageName: string,
544573
methodSignatures: string[],
@@ -757,6 +786,20 @@ export class ModelEditorView extends AbstractWebview<
757786
);
758787
}
759788

789+
private addModeledMethodsFromArray(modeledMethods: ModeledMethod[]) {
790+
const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
791+
792+
for (const modeledMethod of modeledMethods) {
793+
if (!(modeledMethod.signature in modeledMethodsByName)) {
794+
modeledMethodsByName[modeledMethod.signature] = [];
795+
}
796+
797+
modeledMethodsByName[modeledMethod.signature].push(modeledMethod);
798+
}
799+
800+
this.addModeledMethods(modeledMethodsByName);
801+
}
802+
760803
private setModeledMethods(signature: string, methods: ModeledMethod[]) {
761804
this.modelingStore.updateModeledMethods(
762805
this.databaseItem,

‎extensions/ql-vscode/src/model-editor/shared/view-state.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import { QueryLanguage } from "../../common/query-language";
55
export interface ModelEditorViewState {
66
extensionPack: ExtensionPack;
77
language: QueryLanguage;
8-
showFlowGeneration: boolean;
8+
showGenerateButton: boolean;
99
showLlmButton: boolean;
1010
showMultipleModels: boolean;
1111
mode: Mode;

‎extensions/ql-vscode/src/stories/model-editor/LibraryRow.stories.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,7 +219,7 @@ LibraryRow.args = {
219219
modifiedSignatures: new Set(["org.sql2o.Sql2o#Sql2o(String)"]),
220220
inProgressMethods: new Set(),
221221
viewState: createMockModelEditorViewState({
222-
showFlowGeneration: true,
222+
showGenerateButton: true,
223223
showLlmButton: true,
224224
showMultipleModels: true,
225225
}),

‎extensions/ql-vscode/src/stories/model-editor/MethodRow.stories.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -98,7 +98,7 @@ const modeledMethod: ModeledMethod = {
9898
};
9999

100100
const viewState = createMockModelEditorViewState({
101-
showFlowGeneration: true,
101+
showGenerateButton: true,
102102
showLlmButton: true,
103103
showMultipleModels: true,
104104
});

‎extensions/ql-vscode/src/stories/model-editor/ModelEditor.stories.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ ModelEditor.args = {
2828
extensionTargets: {},
2929
dataExtensions: [],
3030
},
31-
showFlowGeneration: true,
31+
showGenerateButton: true,
3232
showLlmButton: true,
3333
showMultipleModels: true,
3434
}),

0 commit comments

Comments
 (0)