diff --git a/extensions/ql-vscode/src/model-editor/languages/models-as-data.ts b/extensions/ql-vscode/src/model-editor/languages/models-as-data.ts index 25ebd21d25f..4fad2426849 100644 --- a/extensions/ql-vscode/src/model-editor/languages/models-as-data.ts +++ b/extensions/ql-vscode/src/model-editor/languages/models-as-data.ts @@ -44,6 +44,7 @@ type ModelsAsDataLanguageModelGeneration = { }; type ModelsAsDataLanguageAccessPathSuggestions = { + queryConstraints: (mode: Mode) => QueryConstraints; parseResults: ( // The results of a single predicate of the query. bqrs: DecodedBqrsChunk, diff --git a/extensions/ql-vscode/src/model-editor/languages/ruby/index.ts b/extensions/ql-vscode/src/model-editor/languages/ruby/index.ts index e1b9d9a5b3e..faa1c6fc73d 100644 --- a/extensions/ql-vscode/src/model-editor/languages/ruby/index.ts +++ b/extensions/ql-vscode/src/model-editor/languages/ruby/index.ts @@ -12,6 +12,7 @@ import { rubyPath, } from "./access-paths"; import { parseAccessPathSuggestionsResults } from "./suggestions"; +import { modeTag } from "../../mode-tag"; export const ruby: ModelsAsDataLanguage = { availableModes: [Mode.Framework], @@ -170,6 +171,10 @@ export const ruby: ModelsAsDataLanguage = { parseResults: parseGenerateModelResults, }, accessPathSuggestions: { + queryConstraints: (mode) => ({ + kind: "table", + "tags contain all": ["modeleditor", "access-paths", modeTag(mode)], + }), parseResults: parseAccessPathSuggestionsResults, }, getArgumentOptions: (method) => { diff --git a/extensions/ql-vscode/src/model-editor/model-editor-view.ts b/extensions/ql-vscode/src/model-editor/model-editor-view.ts index 0680f8f0930..7bb5114d739 100644 --- a/extensions/ql-vscode/src/model-editor/model-editor-view.ts +++ b/extensions/ql-vscode/src/model-editor/model-editor-view.ts @@ -529,6 +529,7 @@ export class ModelEditorView extends AbstractWebview< modelsAsDataLanguage, this.app.logger, ), + queryConstraints: accessPathSuggestions.queryConstraints(mode), cliServer: this.cliServer, queryRunner: this.queryRunner, queryStorageDir: this.queryStorageDir, diff --git a/extensions/ql-vscode/src/model-editor/suggestion-queries.ts b/extensions/ql-vscode/src/model-editor/suggestion-queries.ts index ef96da2d3d3..37ce78fd3a5 100644 --- a/extensions/ql-vscode/src/model-editor/suggestion-queries.ts +++ b/extensions/ql-vscode/src/model-editor/suggestion-queries.ts @@ -1,7 +1,7 @@ import type { CodeQLCliServer } from "../codeql-cli/cli"; import type { Mode } from "./shared/mode"; +import type { QueryConstraints } from "../local-queries"; import { resolveQueriesFromPacks } from "../local-queries"; -import { modeTag } from "./mode-tag"; import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders"; import type { NotificationLogger } from "../common/logging"; import { showAndLogExceptionWithTelemetry } from "../common/logging"; @@ -22,6 +22,7 @@ type RunQueryOptions = { parseResults: ( results: DecodedBqrsChunk, ) => AccessPathSuggestionRow[] | Promise; + queryConstraints: QueryConstraints; cliServer: CodeQLCliServer; queryRunner: QueryRunner; @@ -39,6 +40,7 @@ export async function runSuggestionsQuery( mode: Mode, { parseResults, + queryConstraints, cliServer, queryRunner, logger, @@ -68,6 +70,7 @@ export async function runSuggestionsQuery( cliServer, databaseItem.language, mode, + queryConstraints, ); if (!queryPath) { void showAndLogExceptionWithTelemetry( @@ -141,6 +144,7 @@ export async function runSuggestionsQuery( * @param cliServer The CodeQL CLI server to use. * @param language The language of the query pack to use. * @param mode The mode to resolve the query for. + * @param queryConstraints Constraints to apply to the query. * @param additionalPackNames Additional pack names to search. * @param additionalPackPaths Additional pack paths to search. */ @@ -148,6 +152,7 @@ async function resolveSuggestionsQuery( cliServer: CodeQLCliServer, language: string, mode: Mode, + queryConstraints: QueryConstraints, additionalPackNames: string[] = [], additionalPackPaths: string[] = [], ): Promise { @@ -156,14 +161,7 @@ async function resolveSuggestionsQuery( const queries = await resolveQueriesFromPacks( cliServer, packsToSearch, - { - kind: "table", - "tags contain all": [ - "modeleditor", - "access-path-suggestions", - modeTag(mode), - ], - }, + queryConstraints, additionalPackPaths, ); if (queries.length > 1) { diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts index 336acb2fdfc..0290717ebb1 100644 --- a/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts +++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/model-editor/suggestion-queries.test.ts @@ -1,3 +1,5 @@ +import { load } from "js-yaml"; +import { readFile } from "fs-extra"; import { createMockLogger } from "../../../__mocks__/loggerMock"; import type { DatabaseItem } from "../../../../src/databases/local-databases"; import { DatabaseKind } from "../../../../src/databases/local-databases"; @@ -138,15 +140,21 @@ describe("runSuggestionsQuery", () => { .mockResolvedValueOnce(mockInputSuggestions) .mockResolvedValueOnce(mockOutputSuggestions); + const resolveQueriesInSuite = jest + .fn() + .mockResolvedValue(["/a/b/c/FrameworkModeAccessPathSuggestions.ql"]); + const options = { parseResults, + queryConstraints: { + kind: "table", + "tags all": ["modeleditor", "access-paths", "ruby", "foo"], + }, cliServer: mockedObject({ resolveQlpacks: jest.fn().mockResolvedValue({ "my/extensions": "/a/b/c/", }), - resolveQueriesInSuite: jest - .fn() - .mockResolvedValue(["/a/b/c/FrameworkModeAccessPathSuggestions.ql"]), + resolveQueriesInSuite, packPacklist: jest .fn() .mockResolvedValue([ @@ -205,6 +213,20 @@ describe("runSuggestionsQuery", () => { undefined, undefined, ); + expect(options.cliServer.resolveQueriesInSuite).toHaveBeenCalledTimes(1); + + expect( + load(await readFile(resolveQueriesInSuite.mock.calls[0][0], "utf-8")), + ).toEqual([ + { + from: "codeql/ruby-queries", + include: { + kind: "table", + "tags all": ["modeleditor", "access-paths", "ruby", "foo"], + }, + queries: ".", + }, + ]); expect(options.parseResults).toHaveBeenCalledTimes(2);