Batch policy loads for the Duo workflows GraphQL list

What does this MR do and why?

The Duo session list (getUserWorkflows) no longer spends about 16 SQL queries on every distinct project or group on the page. The policy check for each session reloaded the same settings, root namespace and SAML rows, and re-asked whether the user may use Duo, once per container. The choice to look at: Ai::UserAuthorizable#allowed_to_use now memoizes its answer per request, keyed by user, feature and root namespace.

Design notes

The list field duoWorkflowWorkflows (frontend operation getUserWorkflows, ee/app/assets/javascripts/ai/graphql/get_user_workflow.query.graphql) authorizes every node through Ai::DuoWorkflows::WorkflowPolicy (ee/app/policies/ai/duo_workflows/workflow_policy.rb), which calls can?(:access_duo_agentic_chat, parent) or can?(:duo_workflow, parent) on the session's project or group.

Queries grow per distinct project or group on the page, not per session. Each container cost about 16 queries: project_authorizations max access level for read_project; project_settings for Project#duo_features_enabled; namespaces three times (loading project.group and each container resolving its own root_ancestor object); organizations; namespace_settings twice; namespace_ai_settings; saml_providers via Group#root_saml_provider in the SSO enforcer; and subscription_add_on_purchases five times plus subscription_add_ons once from User#allowed_to_use(:agentic_chat) (or :duo_agent_platform for flows).

Root-level reads repeated because every container built its own root ancestor object, so association caches were never shared even when all projects sit under one group. Namespace#root_ancestor is strong-memoized and Namespace has attr_writer :root_ancestor, which is what the existing Namespaces::Preloaders::*RootAncestorPreloader classes also use.

New class Ai::DuoWorkflows::WorkflowPolicyPreloader (ee/app/models/ai/duo_workflows/workflow_policy_preloader.rb), called from before_connection_authorization in ee/app/graphql/resolvers/ai/duo_workflows/workflows_resolver.rb. It preloads project_setting, organization, project_namespace, group, namespace on projects and organization, namespace_settings on groups; primes max access levels with the existing Preloaders::UserMaxAccessLevelInProjectsPreloader and Preloaders::UserMaxAccessLevelInGroupsPreloader; and loads one Namespace per root id (with namespace_settings, ai_settings, saml_provider preloaded) and assigns it as root_ancestor to every project.namespace, project.group and workflow.namespace object. Dedup is by object_id, since ActiveRecord equality is by id and the same group shows up as several instances.

Why not Preloaders::ProjectPolicyPreloader: its root-ancestor step is a no-op for arrays (it needs an ActiveRecord relation) and it still gives every project its own root object, so root-level reads would repeat.

Ai::UserAuthorizable#allowed_to_use (ee/app/models/concerns/ai/user_authorizable.rb) now wraps its body in Gitlab::SafeRequestStore.fetch keyed by [:duo_allowed_to_use, user id, ai_feature, root_namespace&.id, options]. The result depends only on the user, the root namespace and the options, so a page with N containers under one root asked the same question N times. Outside a request store the method behaves as before. Existing specs in ee/spec/models/concerns/ai/user_authorizable_spec.rb, ee/spec/policies/ai/duo_workflows/workflow_policy_spec.rb, ee/spec/policies/project_policy_spec.rb (Duo sections) and ee/spec/policies/ee/group_policy_spec.rb (Duo sections) pass with the memoization.

Remaining growth: one namespace_settings read per project from CascadingProjectSettingAttribute#locked_ancestor and two per subgroup from CascadingNamespaceSettingAttribute. These come from the cascading duo_features_enabled setting, which has no batch hook. The new request spec ee/spec/requests/api/graphql/ai/duo_workflows/workflows_n_plus_one_spec.rb ignores FROM "namespace_settings" and asserts the rest does not grow with page size, for chat sessions in distinct projects, chat sessions in distinct subgroups, and flow sessions in distinct projects.

The existing N+1 example in ee/spec/requests/api/graphql/ai/duo_workflows/workflows_spec.rb adds sessions in the same project, which is why it never caught this.

Unit spec for the preloader: ee/spec/models/ai/duo_workflows/workflow_policy_preloader_spec.rb.

References

Issue: #629337 (closed) (Graphql N+1 Triage: getUserWorkflows). Epic that describes the triage method: &23614.

Screenshots or screen recordings

No UI change.

How to set up and validate locally

This is backend-only.

  1. On GDK, have a user with a Duo seat and Agentic Chat sessions in at least three different projects under one group. Do this by opening Duo Agentic Chat from three project pages and sending one message each.
  2. Enable the performance bar, or open log/development.log.
  3. Run the query with glab api graphql or curl against http://gdk.test:3000/api/graphql with body { duoWorkflowWorkflows(first: 20) { edges { node { id updatedAt goal aiCatalogItemVersionId agentName archived } } } }.
  4. On master: expect the number of SQL queries to grow by about 16 for every extra project that has a session.
  5. On this branch: expect the count to stay flat, apart from one namespace_settings query per project.

Run the spec with bin/rspec ee/spec/requests/api/graphql/ai/duo_workflows/workflows_n_plus_one_spec.rb.

Measured query counts

8 sessions per page, ActiveRecord::QueryRecorder.new(skip_cached: false), real SaaS setup with Ultimate plan and Duo Pro seat, no policy stubs.

Page layout master this branch
8 sessions in 8 distinct projects 172 63
8 sessions in 8 distinct subgroups 155 65
8 sessions in one project 56 56
Query plans

All new statements are ActiveRecord preloads that look rows up by primary key or unique foreign key for the ids on one page. Plans from Database Lab (gitlab-production-main), ids from the gitlab-org and gitlab-com groups and the gitlab-org/gitlab and gitlab-org/gitlab-foss projects.

SELECT namespaces.* FROM namespaces WHERE namespaces.id IN (9970, 6543)
 Index Scan using namespaces_pkey on public.namespaces  (cost=0.57..6.34 rows=2 width=366) (actual time=2.790..2.811 rows=2 loops=1)
   Index Cond: (namespaces.id = ANY ('{9970,6543}'::integer[]))
   Buffers: shared hit=8 read=2
   I/O Timings: read=2.751 write=0.000
Settings: effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off', random_page_cost = '1.5', seq_page_cost = '4'
Query ID: -6895777607641732531

Time: 5.826 ms
  - planning: 2.950 ms
  - execution: 2.876 ms
    - I/O read: 2.751 ms
    - I/O write: 0.000 ms
SELECT namespace_settings.* FROM namespace_settings WHERE namespace_settings.namespace_id IN (9970, 6543)
 Index Scan using namespace_settings_pkey on public.namespace_settings  (cost=0.56..6.55 rows=2 width=335) (actual time=9.015..16.667 rows=2 loops=1)
   Index Cond: (namespace_settings.namespace_id = ANY ('{9970,6543}'::integer[]))
   Buffers: shared hit=3 read=7
   I/O Timings: read=16.575 write=0.000
Settings: seq_page_cost = '4', effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off', random_page_cost = '1.5'
Query ID: -6835194804201188585

Time: 17.918 ms
  - planning: 1.177 ms
  - execution: 16.741 ms
    - I/O read: 16.575 ms
    - I/O write: 0.000 ms
SELECT namespace_ai_settings.* FROM namespace_ai_settings WHERE namespace_ai_settings.namespace_id IN (9970, 6543)
 Index Scan using index_namespace_ai_settings_on_namespace_id on public.namespace_ai_settings  (cost=0.41..6.39 rows=2 width=90) (actual time=10.325..13.763 rows=2 loops=1)
   Index Cond: (namespace_ai_settings.namespace_id = ANY ('{9970,6543}'::bigint[]))
   Buffers: shared read=5
   I/O Timings: read=13.412 write=0.000
Settings: effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off', random_page_cost = '1.5', seq_page_cost = '4'
Query ID: 1637682880329190557

Time: 14.502 ms
  - planning: 0.680 ms
  - execution: 13.822 ms
    - I/O read: 13.412 ms
    - I/O write: 0.000 ms
SELECT saml_providers.* FROM saml_providers WHERE saml_providers.group_id IN (9970, 6543)
 Index Scan using index_saml_providers_on_group_id on public.saml_providers  (cost=0.28..6.39 rows=2 width=151) (actual time=5.852..5.854 rows=1 loops=1)
   Index Cond: (saml_providers.group_id = ANY ('{9970,6543}'::integer[]))
   Buffers: shared read=3
   I/O Timings: read=5.808 write=0.000
Settings: effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off', random_page_cost = '1.5', seq_page_cost = '4'
Query ID: -8112180808690629733

Time: 6.719 ms
  - planning: 0.808 ms
  - execution: 5.911 ms
    - I/O read: 5.808 ms
    - I/O write: 0.000 ms
SELECT project_settings.* FROM project_settings WHERE project_settings.project_id IN (278964, 13083)
 Index Scan using project_settings_pkey on public.project_settings  (cost=0.56..7.14 rows=2 width=472) (actual time=4.216..4.240 rows=2 loops=1)
   Index Cond: (project_settings.project_id = ANY ('{278964,13083}'::integer[]))
   Buffers: shared hit=7 read=3
   I/O Timings: read=4.176 write=0.000
Settings: random_page_cost = '1.5', seq_page_cost = '4', effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off'
Query ID: 3385664300987609880

Time: 5.615 ms
  - planning: 1.299 ms
  - execution: 4.316 ms
    - I/O read: 4.176 ms
    - I/O write: 0.000 ms
SELECT organizations.* FROM organizations WHERE organizations.id IN (1)
 Index Scan using organizations_pkey on public.organizations  (cost=0.14..3.16 rows=1 width=78) (actual time=5.623..5.626 rows=1 loops=1)
   Index Cond: (organizations.id = 1)
   Buffers: shared hit=3 read=2
   I/O Timings: read=5.582 write=0.000
Settings: effective_cache_size = '472585MB', work_mem = '230MB', jit = 'off', random_page_cost = '1.5', seq_page_cost = '4'
Query ID: 3393604169010413945

Time: 6.417 ms
  - planning: 0.726 ms
  - execution: 5.691 ms
    - I/O read: 5.582 ms
    - I/O write: 0.000 ms

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Edited by Eduardo Bonet

Merge request reports

Loading
Loading