Draft: Add query analyzer to catch legacy checkpoint reads

What this does

Adds a query analyzer that raises when any SQL statement reads the legacy p_duo_workflows_checkpoints table. It raises in development and test, and reports to Sentry in production.

Why

The table p_duo_workflows_checkpoints is being retired. It is replaced by p_duo_workflows_checkpoint_headers and p_duo_workflows_checkpoint_blobs (tracking epic: gitlab-org#23217 (closed)). When the feature flag duo_workflow_write_incremental_only is enabled for a namespace, nothing writes the legacy table any more. Any code that still reads it then silently returns no rows. That is a silent-wrong-data failure, not a crash, so it is easy to miss.

The existing protections do not catch this. A RuboCop cop matches a fixed set of call patterns, so it cannot see association preloads, exists?, count, or new call shapes. Existing runtime tracking instruments only two methods and only records data, it never raises. A new reader gets no signal at all.

How it works

Gitlab::Database::QueryAnalyzers::PreventLegacyDuoWorkflowCheckpointReads lives in ee/lib/gitlab/database/query_analyzers/. It is registered from ee/config/initializers/, because the table is EE-only.

It inspects every SQL statement using the pg_query parse tree. When select_tables contains p_duo_workflows_checkpoints, it raises LegacyCheckpointReadError. Gitlab::Database::QueryAnalyzer#process_sql reports the error to Sentry in production and re-raises it in development and test. Both behaviours come from that existing framework, so the analyzer itself only raises.

Because it reads the parse tree rather than matching Ruby code, it catches shapes the cop cannot: SELECT *, exists? (SELECT 1), COUNT(*), single-column plucks, association preloads, CTEs, subqueries, joins that project only the other table's columns, and INSERT ... SELECT backfills. Pure writes stay exempt, because INSERT, UPDATE, and DELETE never appear in select_tables.

There are two opt-in escape hatches on the model. Ai::DuoWorkflows::Checkpoint.with_legacy_read { } wraps an eager read. Ai::DuoWorkflows::Checkpoint::LegacyRead is a relation extension, used through the legacy_read scope or extending, for relations that run their query later. It overrides exec_queries and count. Association preloads run inside exec_queries, so the extension covers those too.

An ops feature flag, prevent_legacy_duo_workflow_checkpoint_reads, defaults to enabled and gates production reporting only. Development and test always analyze, so a false positive cannot be silenced without anyone noticing.

In specs the analyzer is suppressed by default, because specs that assert on legacy rows directly are legitimate. A spec opts in with legacy_checkpoint_reads: :prevent metadata.

Reads that sit in a flag-off branch and will disappear with the read flag now wrap themselves: the trace.jsonl endpoint in ee/lib/api/ai/duo_workflows/workflows.rb, the internal checkpoint list endpoint in ee/lib/api/ai/duo_workflows/workflows_internal.rb, WorkflowPresenter#first_checkpoint and #latest_checkpoint, Workflow#latest_readable_checkpoint, and CreateCheckpointService#first_checkpoint?. Each wrap is a grep-able marker for the removal work in #611971.

What still reads the legacy table

These readers were never migrated and have no blob path yet, so they return wrong data once duo_workflow_write_incremental_only is on. Wrapping them stops the analyzer from reporting them, so they are listed here as explicit follow-up work in #611971:

  • Workflow#stalled? reports every running session as stalled, because the table is empty.
  • UpdateWorkflowStatusService stops firing the workflow_events_updated GraphQL subscription.
  • WorkflowEventsResolver returns an empty workflowEvents connection.
  • The numeric GET /checkpoints/:checkpoint_id internal endpoint returns 404.

Relationship to the cop MR

This MR targets the branch of !249745 (merged). It should merge after that MR.

It makes two changes to that MR's code. First, it removes the .tap-based runtime tracking. The tracking only instrumented .latest and .earliest, and only when the :workflow association was already loaded. The analyzer sees every statement, so keeping both would report the same read twice.

Second, the cop now treats with_legacy_read and the legacy_read scope as sanctioned markers and skips reads marked either way. Before this change, a read site needed a rubocop:disable comment and a runtime wrapper. Now one marker satisfies both the static check and the runtime check. All the per-site rubocop:disable Gitlab/Ai/AvoidDirectCheckpointTableRead comments are gone.

The internal checkpoint list endpoint now marks its relation with legacy_read instead of wrapping paginate. This matches how the GraphQL resolver marks its relation, and it keeps the header path out of the suppression.

Database review

Danger flagged the new legacy_read scope and the modified with_preloaded_associations scope. Neither changes any SQL. legacy_read is extending(LegacyRead), which attaches a Ruby module to the relation and adds no predicate, join, or ordering. Verified by comparing to_sql with and without it:

-- Ai::DuoWorkflows::Checkpoint.all
SELECT "p_duo_workflows_checkpoints".* FROM "p_duo_workflows_checkpoints"

-- Ai::DuoWorkflows::Checkpoint.legacy_read
SELECT "p_duo_workflows_checkpoints".* FROM "p_duo_workflows_checkpoints"
-- Ai::DuoWorkflows::Workflow.with_preloaded_associations, before and after
SELECT "duo_workflows_workflows".* FROM "duo_workflows_workflows"

Both pairs are byte-identical, so there is no new or modified query to plan. The basic_checkpoints preload that with_preloaded_associations triggers is also unchanged by this MR:

SELECT "p_duo_workflows_checkpoints"."id", "p_duo_workflows_checkpoints"."workflow_id",
       "p_duo_workflows_checkpoints"."created_at", "p_duo_workflows_checkpoints"."thread_ts"
FROM "p_duo_workflows_checkpoints"
WHERE "p_duo_workflows_checkpoints"."workflow_id" IN (1, 2)

What does change is Ruby-level: an extended relation runs its queries with the analyzer suppressed, so a legacy read on that relation is not reported. The overrides cover every route a relation takes to the database (exec_queries, calculate, exists?, pluck), so no allowed read reports by accident and no unmarked read escapes.

Limitations

  • Query analyzers only run inside the Rack middleware and the Sidekiq middleware. Rake tasks and rails console never enable them, so an absence of Sentry reports does not prove there are no readers left.
  • Cached queries are skipped, which is the framework default. If a wrapped read fills the ActiveRecord query cache, a later unwrapped call site with identical SQL is invisible within that request.
  • Queries that name a partition directly, such as gitlab_partitions_dynamic.p_duo_workflows_checkpoints_20260818, are exempt. Application code always reads through the routing table, and partition maintenance has to stay exempt.
  • No changelog entry. The ops flag defaults to enabled, but no user-visible behaviour changes.

Testing

A new analyzer spec covers 19 examples. Nine cover SQL shapes that must raise, five cover shapes that must not, and the rest cover the feature flag, the wrapper, and the error reporting path.

Five new cop examples cover the with_legacy_read block form, the multi-line block form, basic_checkpoints inside a block, the legacy_read chain, and an unmarked read next to a marked one.

New model specs cover both escape hatches, pin each of the four legacy_read query entry points, and confirm the suppression does not leak past the block.

The GraphQL spec ee/spec/requests/api/graphql/ai/duo_workflows_events_spec.rb runs with enforcement enabled. This proves the relation extension covers the real GraphQL pagination path end to end.

All affected suites pass locally, one file at a time: analyzer (19 examples), cop (19), checkpoint model (43), workflow model (430), workflow presenter (26), create checkpoint service (54), update workflow status service (60), internal workflows API (82), workflows API (428), GraphQL workflows (64), GraphQL events (11), and rake task (26). RuboCop is clean on all 17 changed files.

References

Edited by Eduardo Bonet

Merge request reports

Loading
Loading