Wire up CloudConnector::Tokens with Duo Workflow

Matthias Käpplerrequested to merge
526067-use-new-token-types into master

What does this MR do and why?

Part 2 of Add new TokenIssuer types (#526067 - closed).

Introduces CloudConnector::Tokens, which acts as a façade for obtaining a Cloud Connector token. This MR wires up this new API with the CloudConnector::Tokens::* implementations introduced in Add new CloudConnector TokenIssuer types (!185969 - merged) to issue tokens for Duo Workflow.

We chose Duo Workflow as its first consumer, because it is pre-GA, only targets gitlab.com, and because it uses simpler logic than other AI features (for example, it does not take into account customer namespace affinity or add-on seat assignments). This makes it the least risky feature to test the new implementation with.

This change is behind a feature flag: cloud_connector_new_token_impl

References

Database queries

The query we send to establish which Unit Primitives should be available for this request as token scopes has changed.

We current run this query for every AI request.

Before

Querying without namespace filter (relevant for SM/Dedicated):

SELECT
    "subscription_add_on_purchases".*
FROM
    "subscription_add_on_purchases"
    INNER JOIN "subscription_add_ons" "add_on" ON "add_on"."id" = "subscription_add_on_purchases"."subscription_add_on_id"
WHERE
    "add_on"."name" = 3
    AND (started_at IS NULL
        OR started_at <= '2025-04-07')
    AND ('2025-04-07' < expires_on)
-- Formatted by pgFormatter::Beautify
 Hash Join  (cost=3.58..8767.49 rows=30834 width=88) (actual time=3.100..2164.330 rows=61188 loops=1)
   Hash Cond: (subscription_add_on_purchases.subscription_add_on_id = add_on.id)
   Buffers: shared hit=15900 read=4065 dirtied=260
   WAL: records=260 fpi=260 bytes=954168
   I/O Timings: read=2057.067 write=0.000
   ->  Index Scan using index_subscription_addon_purchases_on_expires_on on public.subscription_add_on_purchases  (cost=0.42..8259.42 rows=61668 width=88) (actual time=2.266..2142.269 rows=61673 loops=1)
         Index Cond: (subscription_add_on_purchases.expires_on > '2025-04-07'::date)
         Filter: ((subscription_add_on_purchases.started_at IS NULL) OR (subscription_add_on_purchases.started_at <= '2025-04-07'::date))
         Rows Removed by Filter: 0
         Buffers: shared hit=15897 read=4063 dirtied=260
         WAL: records=260 fpi=260 bytes=954168
         I/O Timings: read=2056.322 write=0.000
   ->  Hash  (cost=3.15..3.15 rows=1 width=8) (actual time=0.806..0.807 rows=1 loops=1)
         Buckets: 1024  Batches: 1  Memory Usage: 9kB
         Buffers: shared hit=3 read=2
         I/O Timings: read=0.744 write=0.000
         ->  Index Scan using index_subscription_add_ons_on_name on public.subscription_add_ons add_on  (cost=0.13..3.15 rows=1 width=8) (actual time=0.794..0.799 rows=1 loops=1)
               Index Cond: (add_on.name = 3)
               Buffers: shared hit=3 read=2
               I/O Timings: read=0.744 write=0.000
Settings: random_page_cost = '1.5', seq_page_cost = '4', work_mem = '100MB', effective_cache_size = '472585MB', jit = 'off'

https://console.postgres.ai/gitlab/gitlab-production-main/sessions/38214/commands/117131

Querying with namespace filter (root group ID; relevant for gitlab.com):

SELECT
    "subscription_add_on_purchases".*
FROM
    "subscription_add_on_purchases"
    INNER JOIN "subscription_add_ons" "add_on" ON "add_on"."id" = "subscription_add_on_purchases"."subscription_add_on_id"
WHERE
    "add_on"."name" = 3
    AND (started_at IS NULL
        OR started_at <= '2025-04-07')
    AND ('2025-04-07' < expires_on)
    AND "subscription_add_on_purchases"."namespace_id" = 9970
-- Formatted by pgFormatter::Beautify
 Nested Loop  (cost=0.55..6.60 rows=1 width=88) (actual time=1.761..1.762 rows=1 loops=1)
   Buffers: shared hit=9 read=3
   I/O Timings: read=1.659 write=0.000
   ->  Index Scan using index_subscription_add_on_purchases_on_namespace_id_add_on_id on public.subscription_add_on_purchases  (cost=0.42..3.44 rows=1 width=88) (actual time=1.737..1.738 rows=1 loops=1)
         Index Cond: (subscription_add_on_purchases.namespace_id = 9970)
         Filter: (((subscription_add_on_purchases.started_at IS NULL) OR (subscription_add_on_purchases.started_at <= '2025-04-07'::date)) AND ('2025-04-07'::date < subscription_add_on_purchases.expires_on))
         Rows Removed by Filter: 0
         Buffers: shared hit=4 read=3
         I/O Timings: read=1.659 write=0.000
   ->  Index Scan using index_subscription_add_ons_on_name on public.subscription_add_ons add_on  (cost=0.13..3.15 rows=1 width=8) (actual time=0.021..0.021 rows=1 loops=1)
         Index Cond: (add_on.name = 3)
         Buffers: shared hit=5
         I/O Timings: read=0.000 write=0.000
Settings: seq_page_cost = '4', work_mem = '100MB', effective_cache_size = '472585MB', jit = 'off', random_page_cost = '1.5'

https://console.postgres.ai/gitlab/gitlab-production-main/sessions/38214/commands/117132

After

Querying without namespace filter (relevant for SM/Dedicated):

SELECT
    "subscription_add_ons".*
FROM
    "subscription_add_ons"
    INNER JOIN "subscription_add_on_purchases" ON "subscription_add_on_purchases"."subscription_add_on_id" = "subscription_add_ons"."id"
WHERE (started_at IS NULL
    OR started_at <= '2025-04-07')
AND ('2025-04-07' < expires_on)
AND "subscription_add_on_purchases"."namespace_id" IS NULL
-- Formatted by pgFormatter::Beautify
  Nested Loop  (cost=0.12..5.73 rows=1 width=57) (actual time=0.002..0.003 rows=0 loops=1)
   Buffers: shared hit=1
   I/O Timings: read=0.000 write=0.000
   ->  Index Scan using index_add_on_purchases_on_add_on_id_and_namespace_id_null on public.subscription_add_on_purchases  (cost=0.12..1.69 rows=1 width=8) (actual time=0.002..0.002 rows=0 loops=1)
         Filter: (((subscription_add_on_purchases.started_at IS NULL) OR (subscription_add_on_purchases.started_at <= '2025-04-07'::date)) AND ('2025-04-07'::date < subscription_add_on_purchases.expires_on))
         Rows Removed by Filter: 0
         Buffers: shared hit=1
         I/O Timings: read=0.000 write=0.000
   ->  Seq Scan on public.subscription_add_ons  (cost=0.00..4.02 rows=2 width=57) (actual time=0.000..0.000 rows=0 loops=0)
         I/O Timings: read=0.000 write=0.000
Settings: effective_cache_size = '472585MB', jit = 'off', random_page_cost = '1.5', seq_page_cost = '4', work_mem = '100MB'

https://postgres.ai/console/gitlab/gitlab-production-main/sessions/38214/commands/117271

Querying with namespace filter (root group ID; relevant for gitlab.com):

SELECT
    "subscription_add_ons".*
FROM
    "subscription_add_ons"
    INNER JOIN "subscription_add_on_purchases" ON "subscription_add_on_purchases"."subscription_add_on_id" = "subscription_add_ons"."id"
WHERE (started_at IS NULL
    OR started_at <= '2025-04-07')
AND ('2025-04-07' < expires_on)
AND "subscription_add_on_purchases"."namespace_id" = 22
-- Formatted by pgFormatter::Beautify
 Nested Loop  (cost=0.42..7.49 rows=1 width=57) (actual time=0.029..0.030 rows=0 loops=1)
   Buffers: shared hit=6
   I/O Timings: read=0.000 write=0.000
   ->  Index Scan using index_subscription_add_on_purchases_on_namespace_id_add_on_id on public.subscription_add_on_purchases  (cost=0.42..3.44 rows=1 width=8) (actual time=0.028..0.029 rows=0 loops=1)
         Index Cond: (subscription_add_on_purchases.namespace_id = 22)
         Filter: (((subscription_add_on_purchases.started_at IS NULL) OR (subscription_add_on_purchases.started_at <= '2025-04-07'::date)) AND ('2025-04-07'::date < subscription_add_on_purchases.expires_on))
         Rows Removed by Filter: 0
         Buffers: shared hit=6
         I/O Timings: read=0.000 write=0.000
   ->  Seq Scan on public.subscription_add_ons  (cost=0.00..4.02 rows=2 width=57) (actual time=0.000..0.000 rows=0 loops=0)
         I/O Timings: read=0.000 write=0.000
Settings: random_page_cost = '1.5', seq_page_cost = '4', work_mem = '100MB', effective_cache_size = '472585MB', jit = 'off'

https://postgres.ai/console/gitlab/gitlab-production-main/sessions/38214/commands/117272

How to set up and validate locally

  1. Set up and run the Duo Workflow service component (and make sure authentication is not bypassed)

  2. Enable the cloud_connector_new_token_impl feature flag

  3. Send the following request:

    curl -v -XPOST -H"Authorization: Bearer <PAT>" gitlab.local:3000/api/v4/ai/duo_workflows/direct_access

This should return a successful JSON response.

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.

Related to #526067 (closed)

Edited by Matthias Käppler

Merge request reports