Skip to content

Commit c584d02

Browse files
starckestarcke
authored
Merge pull request #2734 from github/starcke/automodel-queries-refactor
Starcke/automodel queries refactor
2 parents 1fc7356 + 4fce448 commit c584d02

3 files changed

Lines changed: 181 additions & 154 deletions

File tree

‎extensions/ql-vscode/src/data-extensions-editor/auto-model-codeml-queries.ts‎

Lines changed: 100 additions & 148 deletions
Original file line numberDiff line numberDiff line change
@@ -1,42 +1,21 @@
11
import { CodeQLCliServer, SourceInfo } from "../codeql-cli/cli";
2-
import { QueryRunner } from "../query-server";
2+
import { CoreCompletedQuery, QueryRunner } from "../query-server";
33
import { DatabaseItem } from "../databases/local-databases";
44
import { ProgressCallback } from "../common/vscode/progress";
55
import * as Sarif from "sarif";
66
import { qlpackOfDatabase, resolveQueries } from "../local-queries";
7-
import { extLogger } from "../common/logging/vscode";
87
import { Mode } from "./shared/mode";
9-
import { QlPacksForLanguage } from "../databases/qlpack";
10-
import { createLockFileForStandardQuery } from "../local-queries/standard-queries";
11-
import { CancellationToken, CancellationTokenSource } from "vscode";
128
import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
13-
import { showAndLogExceptionWithTelemetry, TeeLogger } from "../common/logging";
14-
import { QueryResultType } from "../query-server/new-messages";
15-
import { telemetryListener } from "../common/vscode/telemetry";
16-
import { redactableError } from "../common/errors";
179
import { interpretResultsSarif } from "../query-results";
1810
import { join } from "path";
1911
import { assertNever } from "../common/helpers-pure";
2012
import { dir } from "tmp-promise";
2113
import { writeFile, outputFile } from "fs-extra";
2214
import { dump as dumpYaml } from "js-yaml";
2315
import { MethodSignature } from "./external-api-usage";
24-
25-
type AutoModelQueryOptions = {
26-
queryTag: string;
27-
mode: Mode;
28-
cliServer: CodeQLCliServer;
29-
queryRunner: QueryRunner;
30-
databaseItem: DatabaseItem;
31-
qlpack: QlPacksForLanguage;
32-
sourceInfo: SourceInfo | undefined;
33-
additionalPacks: string[];
34-
extensionPacks: string[];
35-
queryStorageDir: string;
36-
37-
progress: ProgressCallback;
38-
token: CancellationToken;
39-
};
16+
import { runQuery } from "../local-queries/run-query";
17+
import { QueryMetadata } from "../common/interface-types";
18+
import { CancellationTokenSource } from "vscode";
4019

4120
function modeTag(mode: Mode): string {
4221
switch (mode) {
@@ -49,108 +28,6 @@ function modeTag(mode: Mode): string {
4928
}
5029
}
5130

52-
async function runAutoModelQuery({
53-
queryTag,
54-
mode,
55-
cliServer,
56-
queryRunner,
57-
databaseItem,
58-
qlpack,
59-
sourceInfo,
60-
additionalPacks,
61-
extensionPacks,
62-
queryStorageDir,
63-
progress,
64-
token,
65-
}: AutoModelQueryOptions): Promise<Sarif.Log | undefined> {
66-
// First, resolve the query that we want to run.
67-
// All queries are tagged like this:
68-
// internal extract automodel <mode> <queryTag>
69-
// Example: internal extract automodel framework-mode candidates
70-
const queries = await resolveQueries(
71-
cliServer,
72-
qlpack,
73-
`Extract automodel ${queryTag}`,
74-
{
75-
kind: "problem",
76-
"tags contain all": ["automodel", modeTag(mode), ...queryTag.split(" ")],
77-
},
78-
);
79-
if (queries.length > 1) {
80-
throw new Error(
81-
`Found multiple auto model queries for ${mode} ${queryTag}. Can't continue`,
82-
);
83-
}
84-
if (queries.length === 0) {
85-
throw new Error(
86-
`Did not found any auto model queries for ${mode} ${queryTag}. Can't continue`,
87-
);
88-
}
89-
90-
const queryPath = queries[0];
91-
const { cleanup: cleanupLockFile } = await createLockFileForStandardQuery(
92-
cliServer,
93-
queryPath,
94-
);
95-
96-
// Get metadata for the query. This is required to interpret the results. We already know the kind is problem
97-
// (because of the constraint in resolveQueries), so we don't need any more checks on the metadata.
98-
const metadata = await cliServer.resolveMetadata(queryPath);
99-
100-
const queryRun = queryRunner.createQueryRun(
101-
databaseItem.databaseUri.fsPath,
102-
{
103-
queryPath,
104-
quickEvalPosition: undefined,
105-
quickEvalCountOnly: false,
106-
},
107-
false,
108-
additionalPacks,
109-
extensionPacks,
110-
queryStorageDir,
111-
undefined,
112-
undefined,
113-
);
114-
115-
const completedQuery = await queryRun.evaluate(
116-
progress,
117-
token,
118-
new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
119-
);
120-
121-
await cleanupLockFile?.();
122-
123-
if (completedQuery.resultType !== QueryResultType.SUCCESS) {
124-
void showAndLogExceptionWithTelemetry(
125-
extLogger,
126-
telemetryListener,
127-
redactableError`Auto-model query ${queryTag} failed: ${
128-
completedQuery.message ?? "No message"
129-
}`,
130-
);
131-
return;
132-
}
133-
134-
const interpretedResultsPath = join(
135-
queryStorageDir,
136-
`interpreted-results-${queryTag.replaceAll(" ", "-")}-${queryRun.id}.sarif`,
137-
);
138-
139-
// eslint-disable-next-line @typescript-eslint/no-unused-vars -- We only need the actual SARIF data, not the extra fields added by SarifInterpretationData
140-
const { t, sortState, ...sarif } = await interpretResultsSarif(
141-
cliServer,
142-
metadata,
143-
{
144-
resultsPath: completedQuery.outputDir.bqrsPath,
145-
interpretedResultsPath,
146-
},
147-
sourceInfo,
148-
["--sarif-add-snippets"],
149-
);
150-
151-
return sarif;
152-
}
153-
15431
type AutoModelQueriesOptions = {
15532
mode: Mode;
15633
candidateMethods: MethodSignature[];
@@ -177,22 +54,13 @@ export async function runAutoModelQueries({
17754
progress,
17855
cancellationTokenSource,
17956
}: AutoModelQueriesOptions): Promise<AutoModelQueriesResult | undefined> {
180-
const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
181-
182-
// CodeQL needs to have access to the database to be able to retrieve the
183-
// snippets from it. The source location prefix is used to determine the
184-
// base path of the database.
185-
const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
57+
// First, resolve the query that we want to run.
58+
const queryPath = await resolveAutomodelQuery(
18659
cliServer,
60+
databaseItem,
61+
"candidates",
62+
mode,
18763
);
188-
const sourceArchiveUri = databaseItem.sourceArchive;
189-
const sourceInfo =
190-
sourceArchiveUri === undefined
191-
? undefined
192-
: {
193-
sourceArchive: sourceArchiveUri.fsPath,
194-
sourceLocationPrefix,
195-
};
19664

19765
// Generate a pack containing the candidate filters
19866
const filterPackDir = await generateCandidateFilterPack(
@@ -205,30 +73,89 @@ export async function runAutoModelQueries({
20573
await cliServer.resolveQlpacks(additionalPacks, true),
20674
);
20775

208-
const candidates = await runAutoModelQuery({
209-
mode,
210-
queryTag: "candidates",
76+
// Run the actual query
77+
const completedQuery = await runQuery({
21178
cliServer,
21279
queryRunner,
21380
databaseItem,
214-
qlpack,
215-
sourceInfo,
81+
queryPath,
82+
queryStorageDir,
21683
additionalPacks,
21784
extensionPacks,
218-
queryStorageDir,
21985
progress,
22086
token: cancellationTokenSource.token,
22187
});
22288

223-
if (!candidates) {
89+
if (!completedQuery) {
22490
return undefined;
22591
}
22692

93+
// Get metadata for the query. This is required to interpret the results. We already know the kind is problem
94+
// (because of the constraint in resolveQueries), so we don't need any more checks on the metadata.
95+
const metadata = await cliServer.resolveMetadata(queryPath);
96+
97+
// CodeQL needs to have access to the database to be able to retrieve the
98+
// snippets from it. The source location prefix is used to determine the
99+
// base path of the database.
100+
const sourceLocationPrefix = await databaseItem.getSourceLocationPrefix(
101+
cliServer,
102+
);
103+
const sourceArchiveUri = databaseItem.sourceArchive;
104+
const sourceInfo =
105+
sourceArchiveUri === undefined
106+
? undefined
107+
: {
108+
sourceArchive: sourceArchiveUri.fsPath,
109+
sourceLocationPrefix,
110+
};
111+
112+
const candidates = await interpretAutomodelResults(
113+
cliServer,
114+
completedQuery,
115+
metadata,
116+
sourceInfo,
117+
);
118+
227119
return {
228120
candidates,
229121
};
230122
}
231123

124+
async function resolveAutomodelQuery(
125+
cliServer: CodeQLCliServer,
126+
databaseItem: DatabaseItem,
127+
queryTag: string,
128+
mode: Mode,
129+
): Promise<string> {
130+
const qlpack = await qlpackOfDatabase(cliServer, databaseItem);
131+
132+
// First, resolve the query that we want to run.
133+
// All queries are tagged like this:
134+
// internal extract automodel <mode> <queryTag>
135+
// Example: internal extract automodel framework-mode candidates
136+
const queries = await resolveQueries(
137+
cliServer,
138+
qlpack,
139+
`Extract automodel ${queryTag}`,
140+
{
141+
kind: "problem",
142+
"tags contain all": ["automodel", modeTag(mode), ...queryTag.split(" ")],
143+
},
144+
);
145+
if (queries.length > 1) {
146+
throw new Error(
147+
`Found multiple auto model queries for ${mode} ${queryTag}. Can't continue`,
148+
);
149+
}
150+
if (queries.length === 0) {
151+
throw new Error(
152+
`Did not found any auto model queries for ${mode} ${queryTag}. Can't continue`,
153+
);
154+
}
155+
156+
return queries[0];
157+
}
158+
232159
/**
233160
* generateCandidateFilterPack will create a temporary extension pack.
234161
* This pack will contain a filter that will restrict the automodel queries
@@ -284,3 +211,28 @@ export async function generateCandidateFilterPack(
284211

285212
return packDir;
286213
}
214+
215+
async function interpretAutomodelResults(
216+
cliServer: CodeQLCliServer,
217+
completedQuery: CoreCompletedQuery,
218+
metadata: QueryMetadata,
219+
sourceInfo: SourceInfo | undefined,
220+
): Promise<Sarif.Log> {
221+
const interpretedResultsPath = join(
222+
completedQuery.outputDir.querySaveDir,
223+
"results.sarif",
224+
);
225+
226+
const { ...sarif } = await interpretResultsSarif(
227+
cliServer,
228+
metadata,
229+
{
230+
resultsPath: completedQuery.outputDir.bqrsPath,
231+
interpretedResultsPath,
232+
},
233+
sourceInfo,
234+
["--sarif-add-snippets"],
235+
);
236+
237+
return sarif;
238+
}

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

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import { CancellationToken } from "vscode";
2+
import { CodeQLCliServer } from "../codeql-cli/cli";
3+
import { ProgressCallback } from "../common/vscode/progress";
4+
import { DatabaseItem } from "../databases/local-databases";
5+
import { CoreCompletedQuery, QueryRunner } from "../query-server";
6+
import { createLockFileForStandardQuery } from "./standard-queries";
7+
import { TeeLogger, showAndLogExceptionWithTelemetry } from "../common/logging";
8+
import { QueryResultType } from "../query-server/new-messages";
9+
import { extLogger } from "../common/logging/vscode";
10+
import { telemetryListener } from "../common/vscode/telemetry";
11+
import { redactableError } from "../common/errors";
12+
import { basename } from "path";
13+
14+
type RunQueryOptions = {
15+
cliServer: CodeQLCliServer;
16+
queryRunner: QueryRunner;
17+
databaseItem: DatabaseItem;
18+
queryPath: string;
19+
queryStorageDir: string;
20+
additionalPacks: string[];
21+
extensionPacks: string[] | undefined;
22+
progress: ProgressCallback;
23+
token: CancellationToken;
24+
};
25+
26+
export async function runQuery({
27+
cliServer,
28+
queryRunner,
29+
databaseItem,
30+
queryPath,
31+
queryStorageDir,
32+
additionalPacks,
33+
extensionPacks,
34+
progress,
35+
token,
36+
}: RunQueryOptions): Promise<CoreCompletedQuery | undefined> {
37+
// Create a lock file for the query. This is required to resolve dependencies and library path for the query.
38+
const { cleanup: cleanupLockFile } = await createLockFileForStandardQuery(
39+
cliServer,
40+
queryPath,
41+
);
42+
43+
// Create a query run to execute
44+
const queryRun = queryRunner.createQueryRun(
45+
databaseItem.databaseUri.fsPath,
46+
{
47+
queryPath,
48+
quickEvalPosition: undefined,
49+
quickEvalCountOnly: false,
50+
},
51+
false,
52+
additionalPacks,
53+
extensionPacks,
54+
queryStorageDir,
55+
undefined,
56+
undefined,
57+
);
58+
59+
const completedQuery = await queryRun.evaluate(
60+
progress,
61+
token,
62+
new TeeLogger(queryRunner.logger, queryRun.outputDir.logPath),
63+
);
64+
65+
await cleanupLockFile?.();
66+
67+
if (completedQuery.resultType !== QueryResultType.SUCCESS) {
68+
void showAndLogExceptionWithTelemetry(
69+
extLogger,
70+
telemetryListener,
71+
redactableError`Failed to run ${basename(queryPath)} query: ${
72+
completedQuery.message ?? "No message"
73+
}`,
74+
);
75+
return;
76+
}
77+
return completedQuery;
78+
}

0 commit comments

Comments
 (0)