Expose Ai::UsageEvent to GraphQL AiUsageData

Robert Huntrequested to merge
549782-expose-ai-usageevent-to-graphql-aiusagedata into master

What does this MR do and why?

This MR exposes the Ai::UsageEvent model to GraphQL under the AiUsageData field.

Within the AiUsageData field it exposes the following properties:

  • ID
  • Timestamp
  • User
  • Event

For now, this endpoint is meant to use PostgreSQL only. Eventually it will support both PSQL and Clickhouse.

References

GraphQL

GraphQL query 
{
  group(fullPath: "gitlab-org") {
    aiUsageData {
    	all {
        nodes {
          id
          timestamp
          user {
            id
            username
          }
          event
        }
        pageInfo {
          endCursor
          hasNextPage
        }
      }
    }
  }
}
GraphQL output 
{
  "data": {
    "group": {
      "aiUsageData": {
        "all": {
          "nodes": [
            {
              "id": "gid://gitlab/Ai::UsageEvent/2",
              "timestamp": "2025-06-30T11:21:56Z",
              "user": {
                "id": "gid://gitlab/User/1",
                "username": "root"
              },
              "event": "TROUBLESHOOT_JOB"
            },
            {
              "id": "gid://gitlab/Ai::UsageEvent/1",
              "timestamp": "2025-07-01T11:21:51Z",
              "user": {
                "id": "gid://gitlab/User/1",
                "username": "root"
              },
              "event": "TROUBLESHOOT_JOB"
            }
          ],
          "pageInfo": {
            "endCursor": "eyJpZCI6IjEifQ",
            "hasNextPage": false
          }
        }
      }
    }
  },
  "correlationId": "01JZAVZDHHVWSPWZBJZYSB2GA0"
}

PSQL Query performance

Query ran

Calling Ai::UsageEventsFinder.new(user, resource: namespace):

SELECT
    *
FROM ( WITH RECURSIVE "array_cte" AS MATERIALIZED (
        SELECT
            "namespaces"."id"
        FROM
            "namespaces"
        WHERE (
            traversal_ids @> (
                '{9970}'
)
)
),
    "recursive_keyset_cte" AS (
(
            SELECT
                NULL::ai_usage_events AS records,
                array_cte_id_array,
                ai_usage_events_timestamp_array,
                ai_usage_events_id_array,
                0::bigint AS count
            FROM (
                SELECT
                    ARRAY_AGG("array_cte"."id") AS array_cte_id_array,
                    ARRAY_AGG("ai_usage_events"."timestamp") AS ai_usage_events_timestamp_array,
                ARRAY_AGG("ai_usage_events"."id") AS ai_usage_events_id_array
            FROM (
                SELECT
                    "array_cte"."id"
                FROM
                    array_cte) array_cte
            LEFT JOIN LATERAL (
                SELECT
                    "ai_usage_events"."timestamp" AS timestamp,
                    "ai_usage_events"."id" AS id
                FROM
                    "ai_usage_events"
                WHERE
                    "ai_usage_events"."timestamp" <= '2025-07-15 19:53:19.629219'
                    AND "ai_usage_events"."namespace_id" = "array_cte"."id"
                ORDER BY
                    "ai_usage_events"."timestamp" DESC,
                    "ai_usage_events"."id" DESC
                LIMIT 1) ai_usage_events ON TRUE
        WHERE
            "ai_usage_events"."timestamp" IS NOT NULL
            AND "ai_usage_events"."id" IS NOT NULL) array_scope_lateral_query
    LIMIT 1)
UNION ALL (
    SELECT
        (
            SELECT
                ai_usage_events
            FROM
                "ai_usage_events"
            WHERE
                "ai_usage_events"."id" = recursive_keyset_cte.ai_usage_events_id_array[position]
            LIMIT 1),
        array_cte_id_array,
        recursive_keyset_cte.ai_usage_events_timestamp_array[:position_query.position - 1] || next_cursor_values.timestamp || recursive_keyset_cte.ai_usage_events_timestamp_array[position_query.position + 1:],
        recursive_keyset_cte.ai_usage_events_id_array[:position_query.position - 1] || next_cursor_values.id || recursive_keyset_cte.ai_usage_events_id_array[position_query.position + 1:],
        recursive_keyset_cte.count + 1
    FROM
        recursive_keyset_cte,
        LATERAL (
            SELECT
                timestamp,
                id,
                position
            FROM
                UNNEST(ai_usage_events_timestamp_array, ai_usage_events_id_array)
                WITH ORDINALITY AS u (timestamp, id, position)
            WHERE
                timestamp IS NOT NULL
                AND id IS NOT NULL
            ORDER BY
                1 DESC,
                2 DESC
            LIMIT 1) AS position_query,
        LATERAL (
            SELECT
                "record"."timestamp",
                "record"."id"
            FROM (
                VALUES (NULL, NULL)) AS nulls
            LEFT JOIN (
                SELECT
                    "ai_usage_events"."timestamp" AS timestamp, "ai_usage_events"."id" AS id
                FROM
                    "ai_usage_events"
                WHERE
                    "ai_usage_events"."timestamp" <= '2025-07-15 19:53:19.629219'
                    AND "ai_usage_events"."namespace_id" = recursive_keyset_cte.array_cte_id_array[position]
                    AND (("ai_usage_events"."timestamp",
                            "ai_usage_events"."id") < (recursive_keyset_cte.ai_usage_events_timestamp_array[position],
                            recursive_keyset_cte.ai_usage_events_id_array[position]))
                ORDER BY
                    "ai_usage_events"."timestamp" DESC,
                    "ai_usage_events"."id" DESC
                LIMIT 1) record ON TRUE
        LIMIT 1) AS next_cursor_values))
SELECT
    (records).*
FROM
    "recursive_keyset_cte" AS "ai_usage_events"
WHERE (count <> 0)) ai_usage_events
LIMIT 100
Pagination
Scenario Result
Default 100 limit default_limit
Overridden limit to 1 limited
Overridden limit to 1 next page limited_next_age

Query plan

https://console.postgres.ai/gitlab/gitlab-production-main/sessions/41528/commands/127573

How to set up and validate locally

Note: If you haven't set up Duo enterprise locally, you can apply this patch to override the check:

Index: ee/app/models/ee/user.rb
===================================================================
diff --git a/ee/app/models/ee/user.rb b/ee/app/models/ee/user.rb
--- a/ee/app/models/ee/user.rb	(revision Staged)
+++ b/ee/app/models/ee/user.rb	(date 1751637798349)
@@ -528,6 +528,8 @@
     end
 
     def assigned_to_duo_enterprise?(container)
+      return true
+
       namespace = ::Gitlab::Saas.feature_available?(:gitlab_duo_saas_only) ? container.root_ancestor : nil
 
       GitlabSubscriptions::AddOnPurchase
  1. In rails console enable the feature flag :unified_ai_events_graphql via http://gdk.test:3000/rails/features
  2. Visit the GraphQL explorer: http://gdk.test:3000/-/graphql-explorer
  3. Enter the GraphQL query above and confirm the output
  4. If you don't have any usage events, you should be able to populate by running FILTER=ai_usage_stats bundle exec rake db:seed_fu

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 #549782 (closed)

Edited by Robert Hunt

Merge request reports