@@ -17,6 +17,10 @@ import { redactableError } from "../common/errors";
1717import { interpretResultsSarif } from "../query-results" ;
1818import { join } from "path" ;
1919import { assertNever } from "../common/helpers-pure" ;
20+ import { dir } from "tmp-promise" ;
21+ import { writeFile , outputFile } from "fs-extra" ;
22+ import { dump as dumpYaml } from "js-yaml" ;
23+ import { MethodSignature } from "./external-api-usage" ;
2024
2125type AutoModelQueryOptions = {
2226 queryTag : string ;
@@ -26,6 +30,7 @@ type AutoModelQueryOptions = {
2630 databaseItem : DatabaseItem ;
2731 qlpack : QlPacksForLanguage ;
2832 sourceInfo : SourceInfo | undefined ;
33+ additionalPacks : string [ ] ;
2934 extensionPacks : string [ ] ;
3035 queryStorageDir : string ;
3136
@@ -52,6 +57,7 @@ async function runAutoModelQuery({
5257 databaseItem,
5358 qlpack,
5459 sourceInfo,
60+ additionalPacks,
5561 extensionPacks,
5662 queryStorageDir,
5763 progress,
@@ -99,7 +105,7 @@ async function runAutoModelQuery({
99105 quickEvalCountOnly : false ,
100106 } ,
101107 false ,
102- getOnDiskWorkspaceFolders ( ) ,
108+ additionalPacks ,
103109 extensionPacks ,
104110 queryStorageDir ,
105111 undefined ,
@@ -147,12 +153,14 @@ async function runAutoModelQuery({
147153
148154type AutoModelQueriesOptions = {
149155 mode : Mode ;
156+ candidateMethods : MethodSignature [ ] ;
150157 cliServer : CodeQLCliServer ;
151158 queryRunner : QueryRunner ;
152159 databaseItem : DatabaseItem ;
153160 queryStorageDir : string ;
154161
155162 progress : ProgressCallback ;
163+ cancellationTokenSource : CancellationTokenSource ;
156164} ;
157165
158166export type AutoModelQueriesResult = {
@@ -161,17 +169,14 @@ export type AutoModelQueriesResult = {
161169
162170export async function runAutoModelQueries ( {
163171 mode,
172+ candidateMethods,
164173 cliServer,
165174 queryRunner,
166175 databaseItem,
167176 queryStorageDir,
168177 progress,
178+ cancellationTokenSource,
169179} : AutoModelQueriesOptions ) : Promise < AutoModelQueriesResult | undefined > {
170- // maxStep for this part is 1500
171- const maxStep = 1500 ;
172-
173- const cancellationTokenSource = new CancellationTokenSource ( ) ;
174-
175180 const qlpack = await qlpackOfDatabase ( cliServer , databaseItem ) ;
176181
177182 // CodeQL needs to have access to the database to be able to retrieve the
@@ -189,17 +194,17 @@ export async function runAutoModelQueries({
189194 sourceLocationPrefix,
190195 } ;
191196
192- const additionalPacks = getOnDiskWorkspaceFolders ( ) ;
197+ // Generate a pack containing the candidate filters
198+ const filterPackDir = await generateCandidateFilterPack (
199+ databaseItem . language ,
200+ candidateMethods ,
201+ ) ;
202+
203+ const additionalPacks = [ ...getOnDiskWorkspaceFolders ( ) , filterPackDir ] ;
193204 const extensionPacks = Object . keys (
194205 await cliServer . resolveQlpacks ( additionalPacks , true ) ,
195206 ) ;
196207
197- progress ( {
198- step : 0 ,
199- maxStep,
200- message : "Finding candidates and examples" ,
201- } ) ;
202-
203208 const candidates = await runAutoModelQuery ( {
204209 mode,
205210 queryTag : "candidates" ,
@@ -208,15 +213,10 @@ export async function runAutoModelQueries({
208213 databaseItem,
209214 qlpack,
210215 sourceInfo,
216+ additionalPacks,
211217 extensionPacks,
212218 queryStorageDir,
213- progress : ( update ) => {
214- progress ( {
215- step : update . step ,
216- maxStep,
217- message : "Finding candidates and examples" ,
218- } ) ;
219- } ,
219+ progress,
220220 token : cancellationTokenSource . token ,
221221 } ) ;
222222
@@ -228,3 +228,59 @@ export async function runAutoModelQueries({
228228 candidates,
229229 } ;
230230}
231+
232+ /**
233+ * generateCandidateFilterPack will create a temporary extension pack.
234+ * This pack will contain a filter that will restrict the automodel queries
235+ * to the specified candidate methods only.
236+ * This is done using the `extensible` predicate "automodelCandidateFilter".
237+ * @param language
238+ * @param candidateMethods
239+ * @returns
240+ */
241+ export async function generateCandidateFilterPack (
242+ language : string ,
243+ candidateMethods : MethodSignature [ ] ,
244+ ) : Promise < string > {
245+ // Pack resides in a temporary directory, to not pollute the workspace.
246+ const packDir = ( await dir ( { unsafeCleanup : true } ) ) . path ;
247+
248+ const syntheticConfigPack = {
249+ name : "codeql/automodel-filter" ,
250+ version : "0.0.0" ,
251+ library : true ,
252+ extensionTargets : {
253+ [ `codeql/${ language } -queries` ] : "*" ,
254+ } ,
255+ dataExtensions : [ "filter.yml" ] ,
256+ } ;
257+
258+ const qlpackFile = join ( packDir , "codeql-pack.yml" ) ;
259+ await outputFile ( qlpackFile , dumpYaml ( syntheticConfigPack ) , "utf8" ) ;
260+
261+ // The predicate has the following defintion:
262+ // extensible predicate automodelCandidateFilter(string package, string type, string name, string signature)
263+ const dataRows = candidateMethods . map ( ( method ) => [
264+ method . packageName ,
265+ method . typeName ,
266+ method . methodName ,
267+ method . methodParameters ,
268+ ] ) ;
269+
270+ const filter = {
271+ extensions : [
272+ {
273+ addsTo : {
274+ pack : `codeql/${ language } -queries` ,
275+ extensible : "automodelCandidateFilter" ,
276+ } ,
277+ data : dataRows ,
278+ } ,
279+ ] ,
280+ } ;
281+
282+ const filterFile = join ( packDir , "filter.yml" ) ;
283+ await writeFile ( filterFile , dumpYaml ( filter ) , "utf8" ) ;
284+
285+ return packDir ;
286+ }
0 commit comments