Add a daily table for GitLab Duo token usage

Andrew Dunnrequested to merge
feature/add-ai-token-daily-usages into master

What does this MR do and why?

Without ClickHouse, a self-managed instance keeps GitLab Duo usage as events for three months, and nothing sums tokens by user, group and model. A platform team on an air-gapped instance needs per-user sums to bill usage back to the programs it hosts. This MR adds the PostgreSQL table and the writer that fills it from the token usage the Duo Chat meter in !259436 (merged) reports. The writer and a per-user index, first planned as a follow-up, are part of this MR.

This MR merges after !260306 (merged), which fixes a detached partition being attached again when the retention period is raised, then lowered, then raised again within a week of the lowering. The setting this MR adds makes that sequence reachable at run time.

  • ai_token_daily_usages has one row per UTC day, organization, namespace chain, user, feature and model. A month is the sum of its days, so a daily limit and a monthly report read the same rows.
  • The chain (traversal_ids) and user ID have no foreign key, so usage stays with the group it was recorded under when a project moves, or when a project or user is deleted. The writer ends a project event's chain at its project namespace and gives usage with no namespace an empty chain.
  • Monthly partitions, so retention drops whole months. The default keeps 1,095 days and duo_token_usage_retention_days accepts 90 to 3,650 days. retain_for now takes a callable, read by the partition manager on every run and documented in the partitioning guide.
  • The previous release's duo_settings schema rejects a row carrying duo_token_usage_retention_days, so this release writes it only into a settings row it creates, keeping an upgraded instance's row savable by the previous release. The form and API parameter are planned for a later release.

Retention path

flowchart LR
  S["duo_token_usage_retention_days"]
  C["retain_for callable"]
  M["Partition manager<br/>every six hours<br/>creates months ahead,<br/>detaches months<br/>past the period"]
  D["Detached partition dropper<br/>daily: drops a month<br/>a week after its detach"]
  P["Empty MINVALUE partition"]
  S ==> C ==> M ==> D
  M -.->|"period raised,<br/>deferred while a detached<br/>one awaits its drop"| P

Raising the period brings no detached month back; it adds only the empty MINVALUE partition above, for rows dated before the oldest attached month, such as a backfill's. Without !260306 (merged), if the period is then lowered and raised again within a week of the lowering, that partition is attached again with those rows, and the dropper keeps it; with it, the partition manager creates that partition empty after the dropper drops the detached table. The floor is set by ai_usage_events, which keeps events for three months: a shorter period would remove only sums the events still hold.

Writer path

When Ai::UsageEvent#store_to_pg accepts a report_duo_chat_token_usage event from Gitlab::Tracking::AiTracking, the writer adds one row to the table's own Redis write buffer for each model in the event's extras.usage that carries token counts. A cron worker drains that buffer every five minutes, as UsageEvents::DumpWriteBufferCronWorker drains the one for ai_usage_events. The worker runs Ai::TokenDailyUsages::DrainBufferService, which calls Ai::TokenDailyUsages::FlushBufferService once per batch until the buffer is empty or 200 seconds have passed.

flowchart LR
  E["AiTracking.track_event<br/>report_duo_chat_token_usage"]
  P["Ai::UsageEvent#store_to_pg<br/>valid events only"]
  W["Ai::TokenDailyUsage.buffer_usage<br/>one row per model,<br/>dated by the event's UTC day"]
  B[("Redis list<br/>ai/token_daily_usage_db_write_buffer")]
  F["Ai::DumpTokenDailyUsageWriteBufferCronWorker<br/>every five minutes"]
  D["Ai::TokenDailyUsages::DrainBufferService<br/>until the buffer is empty<br/>or 200 seconds have passed"]
  S["Ai::TokenDailyUsages::FlushBufferService<br/>one batch per call"]
  T[("ai_token_daily_usages")]
  L[("namespaces<br/>users")]
  E --> P --> W --> B
  F --> D --> S
  B -->|"pops a batch"| S
  S -->|"drops repeated events,<br/>sums each day key,<br/>one upsert_all per batch"| T
  S -->|"current organization of each row:<br/>at most one lookup per table per batch"| L
  S -.->|"upsert fails:<br/>batch goes back"| B
  • Each row is dated by the UTC day of its event's timestamp, so a turn at 23:59:59 UTC and one at 00:00:00 UTC land on two days, whatever the instance's time zone and whenever the flush runs.
  • Each batch is one upsert_all that adds its token and event counts to the day's row. Within a batch, an event the buffer holds twice is counted once, as the ai_usage_events flush drops its repeats.
  • Before the upsert, the flush reads each row's current organization, because an organization transfer can commit between a row being buffered and flushed. A row with a namespace takes its namespace's organization from namespaces, by the last ID of its chain; a row with none takes its user's from users. Each is one primary key lookup per batch, of up to 1,000 IDs. A row keeps its buffered organization when its namespace or user no longer exists. A row left under the old organization has a cost later: after the user is transferred back, a newer row of the same user and day, written under the new organization, moves into the old one and meets the stale row on the unique index, which raises PG::UniqueViolation. A flush that overlaps the transfer's own transaction can still write the old organization.
  • If the upsert fails, nothing is written, the batch goes back to the buffer and the job raises, so the next run writes it once. A row the model rejects, such as a negative count or a day before the retention period, whose partition may already be gone, is dropped so it cannot hold up the rest of the buffer.
  • Only events listed in Ai::TokenDailyUsage::FEATURES_BY_EVENT feed the table, today report_duo_chat_token_usage alone, so usage carried on any other event is not counted. A new meter adds its event there, and a feature value if it is a new feature.
  • The writer keeps its own buffer and worker, so an error in it cannot stall the ai_usage_events flush.
  • No feature flag. The meter from !259436 (merged) checks its own instance setting before it fires report_duo_chat_token_usage, so on an instance with that setting off the buffer stays empty and the table holds no rows.

Organization transfer path

An organization transfer moves a row by its namespace chain, so a group's rows stay together when some of its members move. Organizations::Transfer::GroupsService moves every row recorded under the transferred group, whoever the user. Organizations::Transfer::UsersService moves a transferred user's rows with no namespace and the rows under the user's personal namespace, project namespaces included; a user's rows under any other namespace move when that namespace's group does.

Both transfers call Ai::TokenDailyUsage.transfer_organization, which reads up to 1,000 matching IDs and their dates with no order, moves those rows with the IDs and the distinct dates in the update, so the update plans only the partitions of those months, and reads again until none match. Each update takes its rows out of the source organization, and the method returns at once when the source and target organizations match, so the loop ends. UsersService runs the loop once for a batch of 50 users' rows with no namespace, and once for each of their personal namespaces.

flowchart LR
  subgraph G["Organizations::Transfer::GroupsService, top-level group G"]
    GR["pluck up to 1,000 ids and dates<br/>with no order<br/>organization_id = old<br/>AND traversal_ids under G"]
    GU["update_all<br/>organization_id = new<br/>same predicate<br/>AND id IN ids<br/>AND usage_date IN dates"]
    GD(["done"])
    GR -->|"ids, dates"| GU
    GU -->|"read again"| GR
    GR -->|"no ids"| GD
  end
  subgraph U["Organizations::Transfer::UsersService, each batch of 50 users"]
    UR["pluck up to 1,000 ids and dates<br/>with no order<br/>organization_id = old<br/>AND user_id IN users<br/>AND traversal_ids = '{}'"]
    UU["update_all<br/>organization_id = new<br/>same predicate<br/>AND id IN ids<br/>AND usage_date IN dates"]
    PN["the same loop for each user's<br/>personal namespace N:<br/>traversal_ids under N,<br/>project namespaces included"]
    UD(["done"])
    UR -->|"ids, dates"| UU
    UU -->|"read again"| UR
    UR -->|"no ids"| PN
    PN --> UD
  end
  RO[("the users' rows<br/>under a group's chain")]
  U -.->|"not moved:<br/>they move with their group"| RO
  G ~~~ U

References

How to set up and validate locally

  1. bundle exec rails db:migrate creates the monthly partitions from the retention period back to six months ahead (43 on 7 October).

  2. bundle exec rspec spec/lib/gitlab/database/partitioning/time/base_strategy_spec.rb spec/lib/gitlab/database/partitioning/time/daily_strategy_spec.rb spec/lib/gitlab/database/partitioning/time/weekly_strategy_spec.rb spec/lib/gitlab/database/partitioning/time/monthly_strategy_spec.rb ee/spec/models/ai/token_daily_usage_spec.rb ee/spec/workers/ai/dump_token_daily_usage_write_buffer_cron_worker_spec.rb ee/spec/lib/gitlab/tracking/ai_tracking_spec.rb ee/spec/initializers/1_settings_spec.rb ee/spec/models/application_setting_spec.rb spec/db/schema_spec.rb spec/services/organizations/transfer/users_service_spec.rb ee/spec/services/ee/organizations/transfer/groups_service_spec.rb ee/spec/services/ee/organizations/transfer/users_service_spec.rb ee/spec/services/ai/token_daily_usages/drain_buffer_service_spec.rb ee/spec/services/ai/token_daily_usages/flush_buffer_service_spec.rb passes with ClickHouse running, which most of ai_tracking_spec.rb needs, except the quarantined uses json schema validator, which CI skips; every model and column it lists predates this MR.

  3. With a top-level group and a user, bundle exec rake db:seed_fu FILTER=99_ai_token_daily_usages writes one row, and a second run none.

  4. In a Rails console, with a user who has a personal namespace:

    user = User.first
Gitlab::Tracking::AiTracking.track_event('report_duo_chat_token_usage', user: user, namespace: user.namespace, usage: { 'claude-sonnet-4-5' => { 'input_tokens' => 1200, 'output_tokens' => 300 } })
Ai::DumpTokenDailyUsageWriteBufferCronWorker.new.perform
Ai::TokenDailyUsage.last.slice(:usage_date, :model, :input_tokens, :output_tokens, :events_count)

    The row holds today's UTC date, 1,200 input and 300 output tokens and one event. Running the last three lines again adds to the same row: 2,400, 600 and two events.

Database

Table

  • Every column is NOT NULL.
  • Primary key (id, usage_date).
  • input_tokens, output_tokens and events_count are bigint, default 0, checked >= 0.
  • feature is a smallint enum (chat).
  • model is the AI Gateway's model name, at most 255 characters.
  • Unique (organization_id, traversal_ids, user_id, feature, model, usage_date): the writer's upsert target, and one range in one partition for a group's or project's month or day.
  • The index leads with organization_id, so it serves the cascade without a separate index. The foreign key comes through add_concurrent_partitioned_foreign_key in the same migration; on the empty table its ACCESS EXCLUSIVE lock on organizations lasts milliseconds, under lock retries.
  • The foreign key from organizations is cascade-only by design, so ai_token_daily_usages.organization_id joins spec/support/database/destroy_service_foreign_keys_todo.yml with that reason: a row is an aggregate with nothing outside it to clean up.
  • (user_id, usage_date), for one user's day or month, which a per-user report and a per-user limit will read. It is created with the table, which is empty, so it needs no migration of its own.
  • organization_id is the sharding key, rather than namespace_id, because a row outlives its namespace and some rows have none. The sharding guide keeps group and project data out of tables keyed on organization_id and asks Tenant Scale to approve this choice. A row with a namespace chain takes its namespace's organization and moves with its group on an organization transfer, or with its user when the namespace is the user's personal one; a row with none takes its user's and moves with its user. populate_sharding_key applies the same rule to a row created without an organization.
  • Retention file: drop_partition over the default period, with no pause mechanism, as billable_usage_daily_namespace_aggregates declares. The setting can lengthen the period but not switch retention off; a switch, if the owning group wants one, is the callable returning :ever for a setting value.
  • A regular migration: it creates the table and its indexes in under a second, and the previous release does not use it.

Queries

The writer's upsert is this MR's one new statement besides the organization transfer's reads and updates, by group, by user and by personal namespace, and the flush's two organization lookups; the callable reads cached settings. Locally, with 100,000 tuples on all 31 days and 10 on one (3,100,010 rows in one partition), each statement below reads the unique index:

  • Upsert: 0.24 ms on an existing day, 0.11 ms on a new one.
  • A group's month: 9.7 ms for 15,500 rows; 35 ms on a first run and 16 ms on a second, on a plain copy with rows stored in arrival order.
  • A group's day: 2.3 ms.
  • A project's month: 0.3 ms.
  • The cascade: 0.5 ms across all 43 partitions.

Each local plan, with its SQL:

On a Database Lab clone of GitLab.com, with the table seeded in the same session from September's Duo Chat events, two models each, every plan reads only cached pages:

  • Upsert on an existing day under gitlab-org: 2.9 ms, 2.6 ms of it the LIMIT 1 that picks the row; the writer sends VALUES, so it has no such lookup.
  • gitlab-org's month: 1.7 ms for 442 rows.
  • The gitlab-org/gitlab project's month: 0.9 ms for 246 rows.
Upsert: SQL and plan 
INSERT INTO ai_token_daily_usages (created_at, updated_at, organization_id, user_id, input_tokens,
    output_tokens, events_count, usage_date, feature, traversal_ids, model)
  SELECT now(), now(), organization_id, user_id, 1200, 300, 1, usage_date, feature, traversal_ids, model
  FROM ai_token_daily_usages
  WHERE usage_date >= '2026-09-01' AND usage_date < '2026-10-01'
    AND traversal_ids >= '{9970}' AND traversal_ids < '{9971}'
  LIMIT 1
  ON CONFLICT (organization_id, traversal_ids, user_id, feature, model, usage_date)
  DO UPDATE SET input_tokens = ai_token_daily_usages.input_tokens + EXCLUDED.input_tokens,
    output_tokens = ai_token_daily_usages.output_tokens + EXCLUDED.output_tokens,
    events_count = ai_token_daily_usages.events_count + EXCLUDED.events_count,
    updated_at = EXCLUDED.updated_at;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/166948

The writer's upsert, as it is sent 
INSERT INTO "ai_token_daily_usages" ("organization_id","traversal_ids","user_id","feature","model","usage_date",
    "input_tokens","output_tokens","events_count","created_at","updated_at")
  VALUES (1, '{2}', 1, 0, 'claude-sonnet-4-5', '2026-10-08', 2400, 600, 2, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP),
    (1, '{2}', 1, 0, 'claude-haiku-4-5', '2026-10-08', 80, 20, 1, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
  ON CONFLICT ("organization_id","traversal_ids","user_id","feature","model","usage_date")
  DO UPDATE SET input_tokens = ai_token_daily_usages.input_tokens + excluded.input_tokens,
    output_tokens = ai_token_daily_usages.output_tokens + excluded.output_tokens,
    events_count = ai_token_daily_usages.events_count + excluded.events_count,
    updated_at = excluded.updated_at
  RETURNING "id","usage_date";

This batch held two turns, one with two models and one with one, so the Sonnet row carries two events.

Plans of a one-row form of this statement, on the local seed: https://explain.depesz.com/s/mOsr on an existing day and https://explain.depesz.com/s/6IVm on a new one.

Group's month: SQL and plan 
SELECT user_id, feature, model, SUM(input_tokens) AS input_tokens, SUM(output_tokens) AS output_tokens,
    SUM(events_count) AS events_count
  FROM ai_token_daily_usages
  WHERE usage_date >= '2026-09-01' AND usage_date < '2026-10-01'
    AND organization_id = (SELECT organization_id FROM namespaces WHERE id = 9970)
    AND traversal_ids >= '{9970}' AND traversal_ids < '{9971}'
  GROUP BY user_id, feature, model;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/166949

Project's month: SQL and plan 
SELECT user_id, feature, model, SUM(input_tokens) AS input_tokens, SUM(output_tokens) AS output_tokens,
    SUM(events_count) AS events_count
  FROM ai_token_daily_usages
  WHERE usage_date >= '2026-09-01' AND usage_date < '2026-10-01'
    AND organization_id = (SELECT organization_id FROM namespaces WHERE id = 9970)
    AND traversal_ids >= '{9970,15846663}' AND traversal_ids < '{9970,15846664}'
  GROUP BY user_id, feature, model;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/166950

One user's month and one user's day read the (user_id, usage_date) index. On the same Database Lab clone, with the table seeded from September's Duo Chat events as above, for the user with the most rows (146 in September):

  • One user's month by feature and model: 1.0 ms and 151 pages, against 5.9 ms and 636 pages through the unique index without it.
  • One user's day by model: 0.17 ms and 13 pages, against 3.2 ms and 648 pages through the primary key.
One user's month: SQL and plan

busiest_user holds the user with the most rows and that user's latest day, created in the same session, since a plan cannot show a row's values.

SELECT feature, model, SUM(input_tokens) AS input_tokens, SUM(output_tokens) AS output_tokens,
    SUM(events_count) AS events_count
  FROM ai_token_daily_usages
  WHERE user_id = (SELECT user_id FROM busiest_user)
    AND usage_date >= '2026-09-01' AND usage_date < '2026-10-01'
  GROUP BY feature, model;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/168015

One user's day: SQL and plan 
SELECT model, SUM(input_tokens) AS input_tokens, SUM(output_tokens) AS output_tokens
  FROM ai_token_daily_usages
  WHERE user_id = (SELECT user_id FROM busiest_user)
    AND usage_date = (SELECT usage_date FROM busiest_user)
  GROUP BY model;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/168016

Each transfer's predicate is the unique index's leading columns: organization_id and a range on traversal_ids for a group or a user's personal namespace, and organization_id, traversal_ids = '{}' and the batch's user IDs for 50 users. The statements below are the ones each transfer sends in a spec, with the test's IDs replaced: gitlab-org (9970) in organization 1, :user_ids for the batch's users, :ids and :dates for the IDs and distinct dates the read returned, :namespace_id and :next_namespace_id (the ID plus one) for a personal namespace, and :new_organization_id for the organization the rows move to. Each update keeps the read's predicate beside the IDs, so a row that leaves it after the read is not moved.

On a Database Lab clone, seeded from September's Duo Chat events as above, with each row's organization set as the writer sets it (31,414 rows, one month) and copied over the 11 months before it (371,196 rows, 12 months), each statement ran twice: the first run on whatever the clone had cached, the second warm. The heaviest 50 users' rows (25,410) and the chains of the largest groups and personal namespaces sit under IDs that no account has, so the transfer plans publish no real user's or group's ID besides gitlab-org's. Times are execution times, and an update's include the foreign key check on each row it moves, 5 to 7 ms for 1,000 rows. An update's plan is rolled back, and VACUUM ran between runs. Each plan linked below is the second run. Each transfer update, before and after it carries the batch's dates in usage_date IN, second run first and first run in brackets:

Update Before: ms, partitions planned After: ms, partitions planned
Group, largest group (1,000 of 19,248 rows) 46.2 (108.1), 43 18.9 (18.0), 1
Group, gitlab-org (442 rows, one month) 15.2 (29.4), 43 8.6 (9.4), 1
Group, median group (24 rows over 12 months) 4.8 (4.2), 43 4.0 (4.2), 12
User, heaviest 50 users (1,000 of 25,410 rows) 19.4 (29.9), 43 25.8 (17.9), 1
User, median 50 users (1,000 of 1,200 rows) 26.1 (23.3), 43 32.8 (21.5), 10
Personal namespace, largest (96 rows) new 5.6 (6.2), 12
Personal namespace, median (24 rows) new 3.8 (3.6), 12

The reads return the ID and date of up to 1,000 rows with a LIMIT over the partitions, so they stop at the first partitions that fill the batch and run all 43 only when few rows match:

Read ms Partitions run of 43
Group, largest group (1,000 rows) 1.9 (1.9) 25
Group, gitlab-org (442 rows) 1.5 (1.7) 43
Group, median group (24 rows) 0.8 (1.0) 43
Personal namespace, largest (96 rows) 1.1 (1.2) 43, 31 of them empty
Personal namespace, median (24 rows) 0.7 (3.3) 43
User, heaviest 50 users (1,000 rows) 1.9 (2.0) 25, all rows from one partition
User, median 50 users (1,000 rows) 2.1 (2.2) 34
User, 50 users with no rows outside a namespace 0.8 (1.1) 43, no rows

The flush's two lookups are primary key scans at four to five buffers an ID. They read existing IDs spread evenly across each table's range, with no usage attached. At about 950 IDs, a batch that is nearly full, the first run took 0.9 s for namespaces (909.4 ms, reading 920 pages) and 1.0 s for users (998.4 ms, reading 885 pages), both cold, and the second 7.7 ms and 5.9 ms. At about 95 IDs, both runs took 1.4 to 2.0 ms.

Lookup ms Buffers (second run)
namespaces, 948 IDs 7.7 (909.4 cold) 4,740
namespaces, 96 IDs 2.0 (1.8) 480
users, 932 IDs 5.9 (998.4 cold) 3,713
users, 94 IDs 1.7 (1.4) 376

usage_date IN cuts a group update from 43 partitions to the months holding its batch: one for a batch from a single month, 12 for the median group, whose 24 rows span the year. It takes the largest group's update from 46.2 to 18.9 ms and gitlab-org's from 15.2 to 8.6 ms. The user updates are level within the runs' noise, because 1,000 row writes and their foreign key checks make up most of the time.

Each transfer runs its loop inside the service's transaction, so each read steps over the index entries of the rows already moved. Before the update carried dates, on a clone seeded the same way over a year (371,196 rows) and over three years, the default period (1,113,170 rows), the loop's reads and updates took:

  • The heaviest group: 1.2 s for 19,248 rows at a year and 5.5 s for 57,752 at three years.
  • The heaviest 50 users outside a namespace: 1.6 s for 25,410 rows at a year and 9.0 s for 76,320 at three years.

The reads' slowdown from the rows already moved is 0.07 s and 0.7 s of the group's totals and 0.13 s and 1.4 s of the users'; the updates take most of each total. With a commit after each pass instead, the group read's execution stays near 2 ms throughout.

These loops replace each_batch, whose first query merges every partition's primary key in id order, so a partition with no matching row is read end to end. A group with rows in every month is its best case: there it moved the heaviest group's year in 0.74 s, against the loop's 1.2 s, because each of its updates takes an id range instead of 1,000 listed IDs. A batch with rows in only some months, or with none, as most user batches are, makes its first query read whole months: it filtered a whole month's rows in 8.1 ms for 50 users with no rows outside a namespace, and took 88 ms at a year of rows; for gitlab-org, with rows in one month of twelve, it removed 340,074 rows in 105 ms warm and 181.5 ms cold. !261239 (merged) made the same change for stage events in GroupsService, where each_batch likewise had no index matching both its filter and its order.

Group transfer read: SQL and plan 
SELECT "ai_token_daily_usages"."id", "ai_token_daily_usages"."usage_date" FROM "ai_token_daily_usages"
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND (ai_token_daily_usages.traversal_ids >= '{9970}' AND '{9971}' > ai_token_daily_usages.traversal_ids)
  LIMIT 1000;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170002

Group transfer update: SQL and plan 
UPDATE "ai_token_daily_usages" SET "organization_id" = :new_organization_id
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND (ai_token_daily_usages.traversal_ids >= '{9970}' AND '{9971}' > ai_token_daily_usages.traversal_ids)
    AND "ai_token_daily_usages"."id" IN (:ids)
    AND "ai_token_daily_usages"."usage_date" IN (:dates);

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170028

User transfer read: SQL and plan 
SELECT "ai_token_daily_usages"."id", "ai_token_daily_usages"."usage_date" FROM "ai_token_daily_usages"
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND "ai_token_daily_usages"."user_id" IN (:user_ids)
    AND "ai_token_daily_usages"."traversal_ids" = '{}'
  LIMIT 1000;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170010 for the heaviest 50 users, and https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170070 for 50 users with no rows outside a namespace

User transfer update: SQL and plan 
UPDATE "ai_token_daily_usages" SET "organization_id" = :new_organization_id
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND "ai_token_daily_usages"."user_id" IN (:user_ids)
    AND "ai_token_daily_usages"."traversal_ids" = '{}'
    AND "ai_token_daily_usages"."id" IN (:ids)
    AND "ai_token_daily_usages"."usage_date" IN (:dates);

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170044

Personal namespace transfer read: SQL and plan 
SELECT "ai_token_daily_usages"."id", "ai_token_daily_usages"."usage_date" FROM "ai_token_daily_usages"
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND (ai_token_daily_usages.traversal_ids >= '{:namespace_id}' AND '{:next_namespace_id}' > ai_token_daily_usages.traversal_ids)
  LIMIT 1000;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170006 for the largest personal namespace, and https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170008 for the median

Personal namespace transfer update: SQL and plan 
UPDATE "ai_token_daily_usages" SET "organization_id" = :new_organization_id
  WHERE "ai_token_daily_usages"."organization_id" = 1
    AND (ai_token_daily_usages.traversal_ids >= '{:namespace_id}' AND '{:next_namespace_id}' > ai_token_daily_usages.traversal_ids)
    AND "ai_token_daily_usages"."id" IN (:ids)
    AND "ai_token_daily_usages"."usage_date" IN (:dates);

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170056 for the largest personal namespace

Flush: current organization of a namespace, SQL and plan 
SELECT "namespaces"."id", "namespaces"."organization_id" FROM "namespaces"
  WHERE "namespaces"."id" IN (:namespace_ids);

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170062 for 948 IDs, and https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170064 for 96

Flush: current organization of a user, SQL and plan 
SELECT "users"."id", "users"."organization_id" FROM "users"
  WHERE "users"."id" IN (:user_ids);

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170066 for 932 IDs, and https://postgres.ai/console/gitlab/gitlab-production-main/sessions/65134/commands/170068 for 94

Expected size and access

  • Growth. In September on GitLab.com, by a Database Lab count of ai_usage_events, 8,214 user and namespace pairs used Duo Chat on 15,709 of their days. Assuming two models each, that is about 31,400 rows and 12 MB a month at the 368 bytes a row measured locally before the (user_id, usage_date) index was added: at September's rate, 420 MB over the default period and 1.4 GB over the longest. The primary key, whose entries are as wide as that index's, took about 32 of those bytes. table_size: small. At that rate the table holds about 35 MB at three months, 70 MB at six and 140 MB at a year.
  • Reads and writes. The writer's upsert is the one new write besides the transfer updates above, and each flush also reads namespaces and users by primary key, at most once each. A low-urgency cron job sends it every five minutes, in batches of at most 1,000 buffered rows. At September's rate of about 114 Duo Chat requests an hour (82,373 in the month), and if each request is one turn, a flush sends one small statement of about ten turns. That stays twelve flushes an hour at any horizon, each one statement until five minutes hold more than 1,000 rows. A row changes only when a later turn that day adds to it, or when an organization transfer moves it. Nothing reads the table yet; the planned reads take one day or month of one group or user.
  • Availability risk. Low on GitLab.com and on self-managed instances: the retention period caps growth, at 1.4 GB over the longest period at September's GitLab.com rate, and each planned read takes one day or month of one group or user.
September count: SQL and plan, each figure a node's row count 
WITH chat AS MATERIALIZED (
  SELECT e.user_id, e.namespace_id, (e."timestamp" AT TIME ZONE 'UTC')::date AS usage_date
  FROM ai_usage_events e
  WHERE e.event = 6 AND e."timestamp" >= '2026-09-01 00:00:00+00' AND e."timestamp" < '2026-10-01 00:00:00+00'
)
SELECT
  (SELECT count(*) FROM (SELECT DISTINCT user_id, namespace_id FROM chat) AS p) AS pairs,
  (SELECT count(*) FROM (SELECT DISTINCT user_id, namespace_id, usage_date FROM chat) AS d) AS pair_days,
  (SELECT count(*) FROM (SELECT DISTINCT user_id FROM chat WHERE namespace_id IS NULL) AS n) AS pairs_without_namespace,
  (SELECT count(*) FROM chat) AS requests;

Plan: https://postgres.ai/console/gitlab/gitlab-production-main/sessions/61727/commands/166952

Migration output

PostgreSQL 17: on each database, the dump after rollback matches master's db/structure.sql and the dump after migrating again matches this MR's; ci and sec run the same statements, then lock writes. Advisory-lock and transaction_open? lines are trimmed.

main: == 20261016035646 CreateAiTokenDailyUsages: migrating =========================
main: -- create_table(:ai_token_daily_usages, {:primary_key=>[:id, :usage_date], :options=>"PARTITION BY RANGE (usage_date)", :if_not_exists=>true})
main: -- quote_column_name(:model)
main:    -> 0.0000s
main:    -> 0.0162s
main: -- execute("LOCK TABLE organizations, ai_token_daily_usages IN ACCESS EXCLUSIVE MODE")
main:    -> 0.0007s
main: -- execute("ALTER TABLE ai_token_daily_usages ADD CONSTRAINT fk_rails_f1e54061f3 FOREIGN KEY (organization_id) REFERENCES organizations (id) ON DELETE CASCADE ;")
main:    -> 0.0018s
main: -- execute("ALTER TABLE ai_token_daily_usages VALIDATE CONSTRAINT fk_rails_f1e54061f3;")
main:    -> 0.0006s
main: == 20261016035646 CreateAiTokenDailyUsages: migrated (0.0854s) ================
main: == 20261016035646 CreateAiTokenDailyUsages: reverting =========================
main: -- drop_table(:ai_token_daily_usages)
main:    -> 0.0580s
main: == 20261016035646 CreateAiTokenDailyUsages: reverted (0.0681s) ================

Schema

db/structure.sql holds only this migration's 41 lines; scripts/validate_migration_timestamps and scripts/validate_migration_schema, which db:check-migrations runs, pass locally.

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist.

An AI coding assistant wrote this change, ran the checks above in a disposable development environment, and ran the GitLab.com counts and query plans in Database Lab.

Edited by Andrew Dunn

Merge request reports