perf(governance): resolve customer keys in one bulk query [OM-387]

[gergely-kurucz-konghq]

What

Optimize the governance access query (POST /api/v3/openmeter/governance/query) by resolving customer keys in a single bulk query instead of one DB round trip per key.

Why

resolveCustomers looped GetCustomerByUsageAttribution once per customer key — N serial DB round trips over the full key set (before pagination, so page size didn't help). The e2e benchmark showed this made the customers axis linear; it was the dominant cost at high customer counts.

How

• Add customer.GetCustomersByUsageAttribution(ctx, {Namespace, Keys}) on the customer Service + Adapter: one query matching customer.key IN keys OR subject_key IN keys (mirrors the single-key GetCustomerByUsageAttribution semantics), subjects eager-loaded, soft-deleted rows filtered. Returns []Customer — multi-get-by-set convention, not paginated.
• Rewrite governance resolveCustomers to call it once and reconstruct matched[] / customer-not-found errors in memory from each customer's key + subject keys. Result shape unchanged → sort → paginate → resolveAccess untouched; behavior preserved.

Collapses N customer-resolution queries to 1.

Benchmark (p50, baseline → new): 100×10 −21%, whole C=10 row −17…−29%. High-feature cells flat (100×100 ≈0) — bounded by the still-serial resolveAccess GetAccess loop, a deliberate follow-up.

Summary by CodeRabbit

• New Features
  • Added bulk customer resolution by usage-attribution keys, returning deduplicated matches and supporting lookups via both customer keys and active subject keys.
• Refactor
  • Updated governance access customer resolution to resolve the full key set in one bulk call, then map and report per-key results in-memory.
• Tests
  • Added adapter/service coverage for mixed-key resolution, deduplication, namespace isolation, soft-deleted exclusions, and input validation.
• Documentation
  • Refreshed pagination/phase notes for the new single-bulk-lookup resolution flow.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: c8d48968-8c4c-45cb-ab7e-fb9b77ea51f1

📥 Commits

Reviewing files that changed from the base of the PR and between e816893 and 3fc5025.

📒 Files selected for processing (9)

• openmeter/customer/adapter.go
• openmeter/customer/adapter/customer.go
• openmeter/customer/adapter/customer_test.go
• openmeter/customer/customer.go
• openmeter/customer/service.go
• openmeter/customer/service/customer.go
• openmeter/customer/service/service_test.go
• openmeter/governance/service/service.go
• openmeter/server/server_test.go

💤 Files with no reviewable changes (2)

• openmeter/server/server_test.go
• openmeter/governance/service/service.go

🚧 Files skipped from review as they are similar to previous changes (7)

• openmeter/customer/customer.go
• openmeter/customer/service/service_test.go
• openmeter/customer/adapter.go
• openmeter/customer/adapter/customer_test.go
• openmeter/customer/adapter/customer.go
• openmeter/customer/service.go
• openmeter/customer/service/customer.go

---

📝 Walkthrough

Walkthrough

This PR adds bulk customer resolution by usage attribution keys. A new multi-key input and lookup path are wired through the customer service and adapter, governance now resolves keys in one bulk call, and tests plus mocks cover the new flow.

Changes

Bulk customer resolution by usage attribution keys

Layer / File(s)	Summary
Usage attribution input contract and validation openmeter/customer/customer.go	New GetCustomersByUsageAttributionInput accepts multiple Keys and validates namespace, non-empty keys, and expansions.
Adapter bulk query and tests openmeter/customer/adapter.go, openmeter/customer/adapter/customer.go, openmeter/customer/adapter/customer_test.go	CustomerAdapter adds GetCustomersByUsageAttribution; the adapter resolves customers by key or subject key, converts entities to domain customers, and tests cover matching, deduplication, soft-delete handling, namespace isolation, and validation errors.
Service contract, delegation, and mock openmeter/customer/service.go, openmeter/customer/service/customer.go, openmeter/server/server_test.go, openmeter/customer/service/service_test.go	CustomerService adds the bulk lookup method, validates input before delegating, the server test no-op customer service implements the new method, and service tests cover validation failure for empty keys.
Governance bulk resolution refactor openmeter/governance/service/service.go	resolveCustomers now performs one bulk lookup, maps direct and subject keys in memory, keeps input order, and returns query errors for missing keys.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested reviewers

• GAlexIHU
• rolosp

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 16.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: bulk-resolving customer keys for governance query performance.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ 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/OM-387-governance-performance-optimization

---

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.}

[greptile-apps]

Greptile Summary

This PR optimizes the governance access query endpoint by replacing N serial per-key DB round trips in resolveCustomers with a single bulk GetCustomersByUsageAttribution query. The result shape, sort order, pagination, and resolveAccess phase are all unchanged — only the customer resolution phase is restructured.

• New bulk adapter method: GetCustomersByUsageAttribution issues one query with customer.key IN keys OR subject_key IN keys, eager-loads subjects, and filters soft-deleted rows — mirroring the existing single-key semantics.
• In-memory reconstruction: the governance resolveCustomers function builds a keyToCustomer map from the returned customers' keys and subject keys, then iterates the input keys to reconstruct per-key matches and not-found errors without touching the DB again.
• Validation layering: service-level validation guards the bulk method; adapter-level validation is absent for the new method (noted in a prior review thread).

Confidence Score: 5/5

Safe to merge. The bulk-resolution path is well-isolated, the result contract is unchanged, and the QueryAccessInput.Validate() guard ensures an empty key set never reaches the new code path.

The refactor is narrow: one new adapter method, one service wrapper, and a rewritten resolveCustomers that is straightforwardly correct. The in-memory keyToCustomer map logic handles deduplication the same way the old serial loop did, not-found errors are preserved per-key, and existing test coverage for sort, pagination, and resolveAccess is untouched. The adapter-level validation gap noted in a prior review thread is the only remaining open item, and it is mitigated by service-layer validation.

No files require special attention. The adapter's GetCustomersByUsageAttribution (missing its own input.Validate() call) is the only residual gap, but it was flagged in a prior review thread.

Important Files Changed

Filename	Overview
openmeter/customer/adapter/customer.go	Adds GetCustomersByUsageAttribution adapter method with bulk ent query; missing input.Validate() guard (flagged in previous thread). Logic is otherwise correct.
openmeter/customer/customer.go	Adds GetCustomersByUsageAttributionInput struct and Validate() method; well-formed, all required checks present.
openmeter/customer/service/customer.go	Adds GetCustomersByUsageAttribution service method with validation before delegating to the adapter; correct pattern.
openmeter/governance/service/service.go	Rewrites resolveCustomers to use single bulk lookup; keyToCustomer map construction and dedup logic are correct; doc-comment updated accurately.
openmeter/customer/adapter/customer_test.go	Comprehensive adapter tests cover customer key match, subject key match, dedup, soft-deleted customer/subject, namespace isolation, and unmatched keys.
openmeter/customer/service/service_test.go	Adds service-level validation test ensuring empty-key-set request is rejected as GenericValidationError before hitting the DB.
openmeter/customer/adapter.go	Interface extended with GetCustomersByUsageAttribution signature; straightforward addition.
openmeter/customer/service.go	CustomerService interface extended with GetCustomersByUsageAttribution; straightforward addition.
openmeter/server/server_test.go	NoopCustomerService stub updated to satisfy the extended interface; returns nil, nil which is valid for a noop.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant GovernanceService as governance.service
    participant CustomerService as customer.Service
    participant CustomerAdapter as customer.Adapter
    participant DB

    Client->>GovernanceService: QueryAccess(input)
    GovernanceService->>GovernanceService: input.Validate()
    GovernanceService->>GovernanceService: resolveCustomers(input)

    note over GovernanceService,DB: NEW: single bulk query (was N serial queries)
    GovernanceService->>CustomerService: "GetCustomersByUsageAttribution({Namespace, Keys})"
    CustomerService->>CustomerService: input.Validate()
    CustomerService->>CustomerAdapter: GetCustomersByUsageAttribution(input)
    CustomerAdapter->>DB: "SELECT * FROM customers WHERE namespace=? AND (key IN keys OR HAS_SUBJECTS(subject_key IN keys)) AND deleted_at IS NULL"
    DB-->>CustomerAdapter: []Customer (with subjects eager-loaded)
    CustomerAdapter-->>CustomerService: []Customer
    CustomerService-->>GovernanceService: []Customer

    GovernanceService->>GovernanceService: build keyToCustomer map (in-memory)
    GovernanceService->>GovernanceService: iterate input keys → matched[] + not-found errors (in-memory)
    GovernanceService->>GovernanceService: sort resolved customers (in-memory)
    GovernanceService->>GovernanceService: paginate (in-memory)
    GovernanceService->>GovernanceService: resolveAccess(page) — per-customer GetAccess fan-out (unchanged)
    GovernanceService-->>Client: QueryResult

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Client
    participant GovernanceService as governance.service
    participant CustomerService as customer.Service
    participant CustomerAdapter as customer.Adapter
    participant DB

    Client->>GovernanceService: QueryAccess(input)
    GovernanceService->>GovernanceService: input.Validate()
    GovernanceService->>GovernanceService: resolveCustomers(input)

    note over GovernanceService,DB: NEW: single bulk query (was N serial queries)
    GovernanceService->>CustomerService: "GetCustomersByUsageAttribution({Namespace, Keys})"
    CustomerService->>CustomerService: input.Validate()
    CustomerService->>CustomerAdapter: GetCustomersByUsageAttribution(input)
    CustomerAdapter->>DB: "SELECT * FROM customers WHERE namespace=? AND (key IN keys OR HAS_SUBJECTS(subject_key IN keys)) AND deleted_at IS NULL"
    DB-->>CustomerAdapter: []Customer (with subjects eager-loaded)
    CustomerAdapter-->>CustomerService: []Customer
    CustomerService-->>GovernanceService: []Customer

    GovernanceService->>GovernanceService: build keyToCustomer map (in-memory)
    GovernanceService->>GovernanceService: iterate input keys → matched[] + not-found errors (in-memory)
    GovernanceService->>GovernanceService: sort resolved customers (in-memory)
    GovernanceService->>GovernanceService: paginate (in-memory)
    GovernanceService->>GovernanceService: resolveAccess(page) — per-customer GetAccess fan-out (unchanged)
    GovernanceService-->>Client: QueryResult

Loading

_{Reviews (4): Last reviewed commit: "test(customer): move validation test sce..." | Re-trigger Greptile}