Make query constraints for access paths query configurable

Author: koesie10

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,

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) => {

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,

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<AccessPathSuggestionRow[]>;
+  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,13 +144,15 @@ 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.
  */
 async function resolveSuggestionsQuery(
   cliServer: CodeQLCliServer,
   language: string,
   mode: Mode,
+  queryConstraints: QueryConstraints,
   additionalPackNames: string[] = [],
   additionalPackPaths: string[] = [],
 ): Promise<string | undefined> {
@@ -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) {

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<CodeQLCliServer>({
         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);