fix: optimize customer usage attribution lookup

[turip]

Summary

• replace the cross-table OR in GetCustomerByUsageAttribution with two independently indexable lookup branches combined with UNION ALL
• prefer a direct customer-key match over a subject-key match deterministically
• preserve namespace isolation and existing customer/subject soft-deletion semantics
• add focused adapter coverage and a synthetic PostgreSQL benchmark for both lookup routes

Why

On large customer namespaces, subject-key customer resolution regressed from sub-second behavior to roughly 2-4 seconds per lookup. The old predicate combines a customer-key condition with a correlated subject lookup using a cross-table OR. PostgreSQL can respond by scanning the customer namespace and applying the two alternatives as a filter, even though each lookup is independently indexed.

This lookup is used by entitlement and event-resolution paths, so the per-lookup delay can cascade into API timeouts and worker failures. The new query selects one candidate customer ID through indexed branches, then lets the existing Ent query load that customer normally.

Query plans

Representative local plans below use a synthetic namespace with 100,000 customers and one subject per customer.

Cross-table OR

Limit (actual rows=1 loops=1)
  -> Seq Scan on customers c (actual rows=1 loops=1)
       Filter: (... ((hashed SubPlan 2) OR (key = 'subject-100000')))
       Rows Removed by Filter: 99999
       SubPlan 2
         -> Index Scan using customer_subjects_subject_key on customer_subjects cs
              Index Cond: (subject_key = 'subject-100000')
Planning Time: 0.745 ms
Execution Time: 28.936 ms

The subject lookup uses its index, but the cross-table OR still leaves PostgreSQL scanning and filtering the customer namespace.

UNION ALL candidate lookup

Limit (actual rows=1 loops=1)
  -> Nested Loop (actual rows=1 loops=1)
       -> Limit (actual rows=1 loops=1)
            -> Append (actual rows=1 loops=1)
                 -> Index Scan using customer_namespace_key_active on customers by_key
                      Index Cond: ((namespace = 'benchmark') AND (key = 'subject-100000'))
                 -> Nested Loop (actual rows=1 loops=1)
                      -> Index Scan using customer_subjects_subject_key on customer_subjects cs
                           Index Cond: (subject_key = 'subject-100000')
                      -> Index Scan using customers_pkey on customers by_subject
                           Index Cond: (id = cs.customer_id)
       -> Index Scan using customers_pkey on customers c
            Index Cond: (id = by_key.id)
Planning Time: 0.798 ms
Execution Time: 0.199 ms

Both alternatives enter through their existing indexes and the outer customer fetch uses the selected primary key.

Benchmark

Synthetic local benchmark, 100,000 customers, median of three runs with ten timed iterations per case:

POSTGRES_HOST=127.0.0.1 go test -tags=dynamic ./openmeter/customer/adapter \
  -run '^$' \
  -bench '^BenchmarkCustomerUsageAttributionLookup$' \
  -benchtime=10x \
  -count=3 \
  -benchmem

Lookup route	Cross-table OR	UNION ALL	Improvement
Customer key	35.70 ms/op	3.19 ms/op	11.2x
Subject key	24.46 ms/op	0.95 ms/op	25.7x

Correctness coverage

• customer-key and subject-key matches
• deterministic customer-key precedence when a key also identifies another customer's subject
• customer matching through both its own key and subject key
• namespace isolation
• deleted-customer exclusion
• subject deletion before, exactly at, and after the lookup timestamp
• unchanged not-found behavior

Validation

• make test
• nix develop --impure .#ci -c make lint-go-fast
• go test -tags=dynamic ./openmeter/customer/adapter -count=1
• synthetic benchmark above

Greptile Summary

This PR speeds up customer usage attribution lookup. The main changes are:

• Replaces the cross-table customer/subject OR with a UNION ALL candidate lookup.
• Gives direct customer-key matches deterministic priority over subject-key matches.
• Adds adapter tests for lookup precedence, namespace isolation, and soft-deletion timing.
• Adds a PostgreSQL benchmark and lets the shared DB test helper accept benchmark handles.

Confidence Score: 5/5

This looks safe to merge after checking the helper dependency signature.

• The customer lookup keeps the existing namespace and soft-delete behavior.
• The new precedence behavior is covered by direct adapter tests.
• The only follow-up is confirming the PostgreSQL test helper still matches the external helper API.

openmeter/testutils/pg_driver.go

Important Files Changed

Filename	Overview
openmeter/customer/adapter/customer.go	Reworks the single-key usage attribution lookup into an indexed candidate-ID subquery while keeping namespace and soft-delete filters.
openmeter/customer/adapter/customer_test.go	Adds focused coverage for direct-key lookup, subject-key lookup, precedence, namespace isolation, deleted customers, deleted subjects, and not-found behavior.
openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go	Adds a PostgreSQL-backed benchmark comparing the old cross-table query with the new candidate lookup.
openmeter/testutils/pg_driver.go	Generalizes the PostgreSQL test DB helper to testing.TB so it can be called from benchmarks.

Prompt To Fix All With AI

Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.

---

### Issue 1 of 1
openmeter/testutils/pg_driver.go:132
**Helper Type May Not Match**

`InitPostgresDB` now passes a `testing.TB` interface into `pgtestdb.Custom`. If that dependency's `Custom` function is typed to `*testing.T`, every package that imports this helper will fail to compile even when only normal tests call it; the benchmark support needs a helper path that matches the dependency's accepted test type.

_{Reviews (1): Last reviewed commit: "fix: optimize customer usage attribution..." | Re-trigger Greptile}

Greptile also left 1 inline comment on this PR.

Context used:

• Context used - CLAUDE.md (source)
• Context used - api/spec/AGENTS.md (source)

Summary by CodeRabbit

• Bug Fixes

  • Improved customer lookup accuracy when usage attribution keys match both customer and subject records.
  • Correctly prioritizes direct customer-key matches and prevents duplicate or cross-namespace results.
  • Handles deleted customers and time-based subject deletion rules correctly.

• Performance

  • Optimized usage attribution lookups for faster results at scale.

• Tests

  • Added comprehensive coverage and performance benchmarks for customer lookup scenarios.

[coderabbitai]

Note

Currently processing new changes in this PR. This may take a few minutes, please wait...

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 108d50b1-af46-45dd-9b59-920e9367bff7

📥 Commits

Reviewing files that changed from the base of the PR and between 49b2f18 and ba602ae.

📒 Files selected for processing (4)

• openmeter/customer/adapter/customer.go
• openmeter/customer/adapter/customer_test.go
• openmeter/customer/adapter/customer_usage_attribution_benchmark_test.go
• openmeter/testutils/pg_driver.go

 _____________________________________________________
< Merge conflicts fear my conflict-resolution skills. >
 -----------------------------------------------------
  \
   \   \
        \ /\
        ( )
      .( o ).

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/customer-usage-attribution-benchmark

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted.

Error details

putComment timed out