feat: do not resolve customer for entitlements on each query #3826
Conversation
📝 WalkthroughWalkthroughValidator registration was moved out of NewCustomerService into NewEntitlementRegistry; CustomerService is created earlier and passed into entitlement registry and various entitlement layers. Entitlement model now uses CustomerID instead of embedded Customer; many adapters, services, and wiring updated to be customer-aware. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Init as Init/Wiring
participant CustSvc as CustomerService
participant EntReg as EntitlementRegistry
participant Validator as EntitlementValidator
Init->>CustSvc: NewCustomerService(logger, db, publisher)
activate CustSvc
CustSvc-->>Init: customerService
deactivate CustSvc
Init->>EntReg: NewEntitlementRegistry(..., customerService)
activate EntReg
EntReg->>Validator: NewValidator(entRepo)
activate Validator
Validator-->>EntReg: validator
deactivate Validator
EntReg->>CustSvc: customerService.RegisterRequestValidator(validator)
EntReg-->>Init: entRegistry
deactivate EntReg
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
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. Comment |
A balance check resolves 5-10 times the customer. It's a heavy operation so let's not resolve it that many times.
turip
force-pushed
the
refactor/decrease-db-usage
branch
from
616fcf2 to
48db2fb
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
test/billing/suite.go (1)
146-158:⚠️ Potential issue | 🟠 MajorCustomerService is nil when building the entitlement registry.
GetEntitlementRegistryis called befores.CustomerServiceis initialized (it’s set later), so the registry receives a nil dependency. That can cause panics or missing customer validation in tests. Moving customer setup above the registry creation fixes it.🛠️ Proposed fix (move customer setup before entitlement registry)
@@ - // Entitlement - entitlementRegistry := registrybuilder.GetEntitlementRegistry(registrybuilder.EntitlementOptions{ - DatabaseClient: dbClient, - StreamingConnector: s.MockStreamingConnector, - Logger: slog.Default(), - MeterService: s.MeterAdapter, - CustomerService: s.CustomerService, - Publisher: publisher, - EntitlementsConfiguration: config.EntitlementsConfiguration{ - GracePeriod: datetime.ISODurationString("P1D"), - }, - Locker: locker, - }) + // Customer (needs to exist before entitlements) + customerAdapter, err := customeradapter.New(customeradapter.Config{ + Client: dbClient, + Logger: slog.Default(), + }) + require.NoError(t, err) + + customerService, err := customerservice.New(customerservice.Config{ + Adapter: customerAdapter, + Publisher: publisher, + }) + require.NoError(t, err) + s.CustomerService = customerService + + // Entitlement + entitlementRegistry := registrybuilder.GetEntitlementRegistry(registrybuilder.EntitlementOptions{ + DatabaseClient: dbClient, + StreamingConnector: s.MockStreamingConnector, + Logger: slog.Default(), + MeterService: s.MeterAdapter, + CustomerService: customerService, + Publisher: publisher, + EntitlementsConfiguration: config.EntitlementsConfiguration{ + GracePeriod: datetime.ISODurationString("P1D"), + }, + Locker: locker, + }) @@ - // Customer - customerAdapter, err := customeradapter.New(customeradapter.Config{ - Client: dbClient, - Logger: slog.Default(), - }) - require.NoError(t, err) - - customerService, err := customerservice.New(customerservice.Config{ - Adapter: customerAdapter, - Publisher: publisher, - }) - require.NoError(t, err) - s.CustomerService = customerService + // Customer setup moved above so entitlement registry receives a non-nil CustomerService.openmeter/entitlement/events.go (1)
73-90:⚠️ Potential issue | 🟠 MajorGuard against nil Customer to avoid panics.
ToDomainEntitlementandEventMetadatadereferenceCustomerdirectly; if a caller passes nil (or a legacy payload omits it), this will panic. Consider explicit validation and/or a safe fallback.🛡️ One possible guard
import ( + "errors" "time" @@ func (e entitlementEventV2EntitlementLiteral) Validate() error { + if e.Customer == nil { + return errors.New("customer is required") + } domainEnt := e.ToDomainEntitlement() @@ func (e EntitlementCreatedEventV2) EventMetadata() metadata.EventMetadata { + customerID := "" + if e.Entitlement.Customer != nil { + customerID = e.Entitlement.Customer.ID + } return metadata.EventMetadata{ Source: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityEntitlement, e.Entitlement.ID), - Subject: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityCustomer, e.Entitlement.Customer.ID), + Subject: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityCustomer, customerID), } } @@ func (e EntitlementDeletedEventV2) EventMetadata() metadata.EventMetadata { + customerID := "" + if e.Entitlement.Customer != nil { + customerID = e.Entitlement.Customer.ID + } return metadata.EventMetadata{ Source: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityEntitlement, e.Entitlement.ID), - Subject: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityCustomer, e.Entitlement.Customer.ID), + Subject: metadata.ComposeResourcePath(e.Namespace.ID, metadata.EntityCustomer, customerID), } }Also applies to: 166-170, 197-200
🤖 Fix all issues with AI agents
In `@openmeter/entitlement/driver/v2/entitlement.go`:
- Around line 112-113: The code dereferences r returned from
ParserV2.ToAPIGenericV2 without checking err, which can panic if the function
returns (nil, error); change the call to capture r, err :=
ParserV2.ToAPIGenericV2(&e, cust.ID, cust.Key), then check if err != nil and
immediately return the zero value for the expected return type and err (e.g.,
return APIGenericV2{}, err), otherwise safely dereference r and return *r, nil;
apply this change in the function where ParserV2.ToAPIGenericV2 is invoked to
avoid nil pointer dereference.
In `@openmeter/entitlement/metered/grant_owner_adapter.go`:
- Around line 99-112: The customer deletion check in GrantOwnerAdapter (around
the GetCustomer call) should use the cust.IsDeleted() helper rather than
checking DeletedAt.Before(clock.Now()), so replace the conditional "cust == nil
|| (cust.DeletedAt != nil && cust.DeletedAt.Before(clock.Now()))" with "cust ==
nil || cust.IsDeleted()" in the handler that validates the customer after
customerService.GetCustomer; this aligns behavior with other handlers and
correctly handles the equal-to-now case.
In `@openmeter/subscription/testutils/service.go`:
- Around line 100-103: The test constructs separate instances causing hooks to
attach to a different subject.Service than the one used inside
NewCustomerService; instead instantiate a single Subject service and share it:
create subjectService via NewSubjectService first, then create customerAdapter
with that same subjectService (NewCustomerAdapter(..., subjectService)) and pass
the same subjectService into NewCustomerService (or change NewCustomerService to
accept a subjectService/customerAdapter parameter) so customerAdapter,
customerService and subjectService all reference the same subject.Service
instance and hooks register on the correct object.
🧹 Nitpick comments (8)
openmeter/entitlement/metered/events.go (1)
81-88: Variable name inconsistency:resetEntitlementEventNameV2should beresetEntitlementEventNameV3.The variable is named
V2but it's actually holding the v3 event name (Version: "v3"). This is confusing for future maintainers.Also, the interface check on line 82 verifies
EntitlementResetEvent(v1) implementsmarshaler.Event, but there's no check forEntitlementResetEventV3.♻️ Suggested fix
var ( - _ marshaler.Event = EntitlementResetEvent{} + _ marshaler.Event = EntitlementResetEvent{} + _ marshaler.Event = EntitlementResetEventV3{} - resetEntitlementEventNameV2 = metadata.GetEventName(metadata.EventType{ + resetEntitlementEventNameV3 = metadata.GetEventName(metadata.EventType{ Subsystem: EventSubsystem, Name: "entitlement.reset", Version: "v3", }) )And update the reference in
EventName():func (e EntitlementResetEventV3) EventName() string { - return resetEntitlementEventNameV2 + return resetEntitlementEventNameV3 }openmeter/entitlement/metered/reset.go (1)
56-65: Consider usingclock.Now()instead oftime.Now()for testability.Line 64 uses
time.Now()forResetRequestedAt, but elsewhere in the codebase (and in tests),clock.Now()is used to enable time manipulation in tests. This could make it harder to write deterministic tests for reset event timestamps.♻️ Suggested fix
CustomerID: ent.CustomerID, ResetAt: params.At, RetainAnchor: params.RetainAnchor, - ResetRequestedAt: time.Now(), + ResetRequestedAt: clock.Now(), }openmeter/entitlement/driver/entitlement.go (1)
276-286: Consider reordering nil check before IsDeleted check.The current code checks
cust.IsDeleted()at line 276 before checking ifcust == nilat line 282. Ifcustis nil, callingcust.IsDeleted()would panic. However, looking more carefully, the nil check for thecust != nilcondition is part of the compound condition on line 276, so this is actually safe.Actually, wait - line 276 has
if cust != nil && cust.IsDeleted(), so the nil case falls through to lines 282-286. This is correct but a bit scattered. Consider consolidating:♻️ Suggested consolidation for clarity
+ if cust == nil { + return GetEntitlementsOfSubjectHandlerResponse{}, models.NewGenericPreConditionFailedError( + fmt.Errorf("customer not found [namespace=%s subject.id=%s]", id.Namespace, id.SubjectIdOrKey), + ) + } + - if cust != nil && cust.IsDeleted() { + if cust.IsDeleted() { return GetEntitlementsOfSubjectHandlerResponse{}, models.NewGenericPreConditionFailedError( fmt.Errorf("customer is deleted [namespace=%s customer.id=%s]", cust.Namespace, cust.ID), ) } - - if cust == nil { - return GetEntitlementsOfSubjectHandlerResponse{}, models.NewGenericPreConditionFailedError( - fmt.Errorf("customer not found [namespace=%s subject.id=%s]", id.Namespace, id.SubjectIdOrKey), - ) - }openmeter/entitlement/uniqueness.go (1)
37-41: Consider updating the error message for clarity.The grouping is now by
CustomerID, but the error message still says "entitlements must belong to the same subject". Since "subject" and "customer" appear to be distinct concepts in this codebase, it might be worth updating the message to say "entitlements must belong to the same customer" to avoid confusion during debugging.💡 Suggested change
if grouped := lo.GroupBy(ents, func(e Entitlement) string { return e.CustomerID }); len(grouped) > 1 { keys := lo.Keys(grouped) slices.Sort(keys) - return fmt.Errorf("entitlements must belong to the same subject, found %v", keys) + return fmt.Errorf("entitlements must belong to the same customer, found %v", keys) }openmeter/entitlement/balanceworker/worker.go (1)
287-298: Comment mentions V2 but handler processes V3.The comment on line 287 says "Metered entitlement reset event v2" but the handler now processes
EntitlementResetEventV3. Consider updating the comment to match.📝 Suggested comment update
- // Metered entitlement reset event v2 + // Metered entitlement reset event v3 grouphandler.NewGroupEventHandler(func(ctx context.Context, event *meteredentitlement.EntitlementResetEventV3) error {openmeter/entitlement/service/scheduling.go (1)
28-130: Consider deferring the customer lookup until after feature/uniqueness checks.
This avoids a DB hit when validation fails early (e.g., missing feature or uniqueness conflict).♻️ Possible refactor
- customer, err := c.customerService.GetCustomer(ctx, customer.GetCustomerInput{ - CustomerID: &customer.CustomerID{ - Namespace: input.Namespace, - ID: input.UsageAttribution.ID, - }, - }) - if err != nil { - return nil, err - } + // ... after feature lookup + uniqueness validation ... + customer, err := c.customerService.GetCustomer(ctx, customer.GetCustomerInput{ + CustomerID: &customer.CustomerID{ + Namespace: input.Namespace, + ID: input.UsageAttribution.ID, + }, + }) + if err != nil { + return nil, err + }As per coding guidelines: Performance should be a priority in critical code paths. Anything related to event ingestion, message processing, database operations (regardless of database) should be vetted for potential performance bottlenecks.
openmeter/entitlement/driver/parser.go (2)
35-41: Consider extracting the repeated subject key resolution logic.The same pattern for extracting
subjKeyfrome.Customer.UsageAttributionappears three times (inToMetered,ToStatic, andToBoolean). Extracting this to a small helper would reduce duplication and make future changes easier.♻️ Suggested helper extraction
// Add this helper function somewhere in the file func getSubjectKey(c customer.Customer) string { if c.UsageAttribution == nil { return "" } key, err := c.UsageAttribution.GetFirstSubjectKey() if err != nil { return "" } return key }Then each method simply calls:
subjKey := getSubjectKey(e.Customer)Also applies to: 80-86, 119-125
52-52: Placeholder comment forIsUnlimitedfield.There's a
// implementcomment here indicating this is hardcoded tofalse. If unlimited entitlements are planned, it'd be good to track this as a follow-up task so it doesn't get lost.Would you like me to open an issue to track implementing the
IsUnlimitedlogic?
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openmeter/entitlement/driver/entitlement.go (1)
103-115:⚠️ Potential issue | 🟠 MajorGuard against
nilcustomer before dereferencing.If
resolveCustomerFromSubjectreturns(nil, nil), bothcust.GetUsageAttribution()and*custwill panic. Please add a nil check and return a clear precondition error in both Create and Override flows.🔧 Suggested fix
func (h *entitlementHandler) CreateEntitlement() CreateEntitlementHandler { return httptransport.NewHandlerWithArgs[CreateEntitlementHandlerRequest, CreateEntitlementHandlerResponse, string]( @@ func(ctx context.Context, request CreateEntitlementHandlerRequest) (CreateEntitlementHandlerResponse, error) { cust, err := h.resolveCustomerFromSubject(ctx, request.Namespace, request.SubjectIdOrKey) if err != nil { return nil, err } + if cust == nil { + return nil, models.NewGenericPreConditionFailedError( + fmt.Errorf("customer not found [namespace=%s subject.id=%s]", request.Namespace, request.SubjectIdOrKey), + ) + } request.Inputs.UsageAttribution = cust.GetUsageAttribution() @@ func (h *entitlementHandler) OverrideEntitlement() OverrideEntitlementHandler { return httptransport.NewHandlerWithArgs[OverrideEntitlementHandlerRequest, OverrideEntitlementHandlerResponse, OverrideEntitlementHandlerParams]( @@ func(ctx context.Context, request OverrideEntitlementHandlerRequest) (OverrideEntitlementHandlerResponse, error) { cust, err := h.resolveCustomerFromSubject(ctx, request.Inputs.Namespace, request.SubjectIdOrKey) if err != nil { return nil, err } + if cust == nil { + return nil, models.NewGenericPreConditionFailedError( + fmt.Errorf("customer not found [namespace=%s subject.id=%s]", request.Inputs.Namespace, request.SubjectIdOrKey), + ) + } request.Inputs.UsageAttribution = cust.GetUsageAttribution()Also applies to: 167-179
🧹 Nitpick comments (3)
openmeter/entitlement/driver/v2/entitlement.go (1)
151-156: Variable shadows the imported package name.The local variable
entitlementat line 151 shadows the importedentitlementpackage (line 13). It works here since the package isn't used after this point in the scope, but it can trip up readers or future edits. Consider renaming to something likeentorresultfor clarity.♻️ Suggested rename
- entitlement, err := h.connector.GetEntitlementWithCustomer(ctx, request.Namespace, request.EntitlementId) + ent, err := h.connector.GetEntitlementWithCustomer(ctx, request.Namespace, request.EntitlementId) if err != nil { return nil, err } - return ParserV2.ToAPIGenericV2(&entitlement.Entitlement, entitlement.Customer.ID, entitlement.Customer.Key) + return ParserV2.ToAPIGenericV2(&ent.Entitlement, ent.Customer.ID, ent.Customer.Key)openmeter/entitlement/metered/reset.go (1)
57-65: Guard against missing CustomerID before publishing.If any legacy entitlements lack
CustomerID, the v3 event validation will fail after the reset work, and the call returns an error. Can you confirm the invariant (or backfill)? If not guaranteed, consider an early guard for clearer failure.🔧 Possible early validation
+ if ent.CustomerID == "" { + return nil, errors.New("customerID must be set for metered entitlement reset") + } + event := EntitlementResetEventV3{ EntitlementID: entitlementID.ID, Namespace: eventmodels.NamespaceID{ ID: entitlementID.Namespace, }, CustomerID: ent.CustomerID, ResetAt: params.At, RetainAnchor: params.RetainAnchor, ResetRequestedAt: clock.Now(), }openmeter/entitlement/metered/events.go (1)
81-88: Add a compile-time interface assertion for the v3 event.Small maintainability win: you already assert
EntitlementResetEventimplementsmarshaler.Event; adding the same forEntitlementResetEventV3makes future refactors safer.♻️ Suggested tweak
var ( _ marshaler.Event = EntitlementResetEvent{} + _ marshaler.Event = EntitlementResetEventV3{} resetEntitlementEventNameV3 = metadata.GetEventName(metadata.EventType{ Subsystem: EventSubsystem, Name: "entitlement.reset", Version: "v3", }) )
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@test/billing/suite.go`:
- Around line 169-181: The EntitlementOptions passed into
registrybuilder.GetEntitlementRegistry is missing a Tracer; add Tracer:
noop.NewTracerProvider().Tracer("test_env") to the EntitlementOptions literal so
NewEntitlementGrantOwnerAdapter, creditConnector, and
NewMeteredEntitlementConnector receive a non-nil tracer; locate the
EntitlementOptions block used in registrybuilder.GetEntitlementRegistry and add
the Tracer field consistent with other test setup code that uses
noop.NewTracerProvider().Tracer("test_env").
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
openmeter/entitlement/metered/events.go (1)
17-17:⚠️ Potential issue | 🟡 MinorFix the deprecation comment to point to V3 instead of the nonexistent V2.
The deprecation comment at line 17 says to use
EntitlementResetEventV2, but that type was completely removed—V2 doesn't exist anywhere in the codebase. V3 is what's actually being used now. The comment should point toEntitlementResetEventV3instead to avoid sending developers down a confusing dead-end during migrations.
GAlexIHU
left a comment
GAlexIHU
left a comment
There was a problem hiding this comment.
I get why this change is wanted and i'm ok with it being merged but I don't really like the way it's done.
Some thoughts on possible approaches:
- If we're removing customer from entitlement then we could just use the customer service separately every place customer info is needed
e.g. roughlt
ents, err := s.entitlementService.ListEntitlements(...)
custs, err := s.customerService.ListCustomers(mapEntsToCustIds(ents))
out, err := matchCustomersToEnts(ents, custs)-
If we're not removing it then we already have the expand pattern and optional types in the codebase, we could populate the field conditionally.
-
If this is primarily a concern for value checks IMO we should just rewrite that flow so it pre-fetches everything approximately once (though I recognize that's significantly more effort than this)
Also, though I recognize that from OMC's perspective this could save some costs, will we actually materialize that?
These are just thoughts / conversation topics, I'm fine with the change as is (except that one comment), I don't expect any of this to change
| } | ||
|
|
||
| // Create and register the entitlement validator | ||
| validator, err := entitlementvalidator.NewValidator(entRegistry.EntitlementRepo) |
There was a problem hiding this comment.
i'm not sure if there's a service where customer svc is used but not entitlement, but if there is, customer svc should still have the entitlement validator, which it won't after this change
Just from propsperity's sake: |
Overview
A balance check resolves 7 times the customer. It's a heavy operation so let's not resolve it that many times.
This will also shave off 50ms from the calculation time.
This patch decreases that number by at least 5.
Notes for reviewer
Summary by CodeRabbit
New Features
Refactor
Tests
✏️ Tip: You can customize this high-level summary in your review settings.