fix: entitlement query for customer meters #3758
Conversation
📝 WalkthroughWalkthroughAdds new DB indexes on customers and features, and refactors entitlement listing to accept a single ingest-event filter and perform a consolidated entitlement query instead of iterating over multiple filters. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Handler as BalanceWorker.IngestHandler
participant Repo as BalanceWorkerRepository
participant DB as Database
rect rgb(230,245,255)
Handler->>Repo: ListEntitlementsAffectedByIngestEvents(filter)
note right of Repo: Resolve IDs once
end
rect rgb(255,250,230)
Repo->>DB: Query Customers WHERE namespace & subject
DB-->>Repo: customerIDs
Repo->>DB: Query Features WHERE namespace AND meter_slugs (archived_at IS NULL)
DB-->>Repo: featureIDs
Repo->>DB: Query Entitlements WHERE namespace AND customerID IN (...) AND featureID IN (...)
DB-->>Repo: entitlements
end
rect rgb(230,245,255)
Repo-->>Handler: []ListAffectedEntitlementsResponse (namespace, entitlement_id, created_at, deleted_at, active_from, active_to)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI Agents
In @openmeter/entitlement/adapter/entitlement.go:
- Around line 276-295: The features query result error is not checked and gets
overwritten by the subsequent entitlements query; after calling
repo.db.Feature.Query().Where(...).All(ctx) check if err != nil and return nil,
err before proceeding to build the entitlements query (ensure you still use the
features variable for db_entitlement.FeatureIDIn and customers for
db_entitlement.CustomerIDIn); this prevents losing the original features query
error.
🧹 Nitpick comments (1)
openmeter/entitlement/adapter/entitlement.go (1)
259-295: Consider short-circuiting when no customers or features are found.If either
customersorfeaturesis empty, callingCustomerIDIn()orFeatureIDIn()with zero arguments could lead to unexpected behavior depending on how the ORM generates SQL. Returning early with an empty result would be more predictable and efficient.🔎 Proposed fix
customers, err := repo.db.Customer.Query().Where( customerdb.Namespace(eventFilter.Namespace), customerNotDeletedAt(clock.Now()), customerdb.Or( customerdb.Key(eventFilter.EventSubject), customerdb.HasSubjectsWith( customersubjectsdb.SubjectKey(eventFilter.EventSubject), customersubjectsdb.DeletedAtIsNil(), ), ), ).All(ctx) if err != nil { return nil, fmt.Errorf("failed to query customers: %w", err) } +if len(customers) == 0 { + return result, nil +} // resolve feature IDs features, err := repo.db.Feature.Query().Where( db_feature.Namespace(eventFilter.Namespace), db_feature.MeterSlugIn(eventFilter.MeterSlugs...), db_feature.ArchivedAtIsNil(), ).All(ctx) +if err != nil { + return nil, fmt.Errorf("failed to query features: %w", err) +} +if len(features) == 0 { + return result, nil +}
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (2)
openmeter/ent/db/migrate/schema.gois excluded by!**/ent/db/**tools/migrate/migrations/atlas.sumis excluded by!**/*.sum,!**/*.sum
📒 Files selected for processing (8)
openmeter/ent/schema/customer.goopenmeter/ent/schema/feature.goopenmeter/entitlement/adapter/entitlement.goopenmeter/entitlement/balanceworker/ingesthandler.goopenmeter/entitlement/balanceworker/repository.gotest/entitlement/regression/scenario_test.gotools/migrate/migrations/20260106184018_balance_worker_idx_improvements.down.sqltools/migrate/migrations/20260106184018_balance_worker_idx_improvements.up.sql
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go
⚙️ CodeRabbit configuration file
**/*.go: In general when reviewing the Golang code make readability and maintainability a priority, even potentially suggest restructuring the code to improve them.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.
Files:
openmeter/entitlement/adapter/entitlement.goopenmeter/entitlement/balanceworker/ingesthandler.goopenmeter/ent/schema/feature.goopenmeter/entitlement/balanceworker/repository.gotest/entitlement/regression/scenario_test.goopenmeter/ent/schema/customer.go
**/*_test.go
⚙️ CodeRabbit configuration file
**/*_test.go: Make sure the tests are comprehensive and cover the changes. Keep a strong focus on unit tests and in-code integration tests.
When appropriate, recommend e2e tests for critical changes.
Files:
test/entitlement/regression/scenario_test.go
🧠 Learnings (1)
📚 Learning: 2025-10-09T13:59:12.012Z
Learnt from: chrisgacsal
Repo: openmeterio/openmeter PR: 3486
File: openmeter/ingest/kafkaingest/serializer/serializer.go:105-107
Timestamp: 2025-10-09T13:59:12.012Z
Learning: In OpenMeter, the CloudEvents `subject` field is mandatory for the application's business logic, even though it's optional in the CloudEvents specification. The `ValidateKafkaPayloadToCloudEvent` function in `openmeter/ingest/kafkaingest/serializer/serializer.go` intentionally enforces this requirement.
Applied to files:
test/entitlement/regression/scenario_test.go
🧬 Code graph analysis (4)
openmeter/entitlement/adapter/entitlement.go (2)
openmeter/entitlement/balanceworker/repository.go (2)
IngestEventQueryFilter(16-20)ListAffectedEntitlementsResponse(22-29)pkg/framework/entutils/transaction.go (1)
TransactingRepo(199-221)
openmeter/entitlement/balanceworker/ingesthandler.go (2)
openmeter/entitlement/balanceworker/repository.go (1)
IngestEventQueryFilter(16-20)openmeter/subject/subject.go (1)
SubjectKey(34-36)
openmeter/ent/schema/feature.go (1)
pkg/models/annotation.go (1)
Annotations(5-5)
test/entitlement/regression/scenario_test.go (2)
openmeter/entitlement/balanceworker/repository.go (1)
IngestEventQueryFilter(16-20)app/common/namespace.go (1)
Namespace(13-15)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
- GitHub Check: Lint
- GitHub Check: Artifacts / Benthos Collector Container image
- GitHub Check: Artifacts / Container image
- GitHub Check: Test
- GitHub Check: Code Generators
- GitHub Check: Build
- GitHub Check: Migration Checks
- GitHub Check: Repository Scan
- GitHub Check: Analyze (go)
🔇 Additional comments (9)
tools/migrate/migrations/20260106184018_balance_worker_idx_improvements.up.sql (2)
3-4: LGTM! The feature index looks well-designed.The partial index on
(namespace, meter_slug)with theWHERE archived_at IS NULLcondition is a solid choice. It'll speed up queries that filter by namespace and meter slug while excluding archived features, and the smaller index size (due to the partial condition) is a nice bonus.
1-2: Index is already defined in the schema—no new concerns here.This index already exists in
openmeter/ent/schema/customer.go(line 49) and follows the same pattern asCustomerSubjectsandSubjectentities. Most queries use the unique index on(namespace, key) WHERE deleted_at IS NULL, which covers the typical case of looking up active customers. The broader index on(namespace, key, deleted_at)is intentionally included to support soft-delete synchronization logic and range scans on deleted records—it's consistent with how the codebase handles soft deletes across related entities.tools/migrate/migrations/20260106184018_balance_worker_idx_improvements.down.sql (1)
1-4: LGTM! Clean rollback migration.The down migration correctly drops both indexes in reverse order, providing a clean rollback path.
openmeter/ent/schema/feature.go (1)
46-49: LGTM! Index definition looks solid.The new composite index on
(namespace, meter_slug)with the partial filter for active features aligns perfectly with the SQL migration and follows the same pattern as the existing indexes in this schema. This should help with queries that look up features by meter slug within a namespace.openmeter/ent/schema/customer.go (1)
49-49: Same index utility concern as the SQL migration.This composite index on
(namespace, key, deleted_at)matches the SQL migration, but as I noted in the migration file review, its effectiveness depends on the actual query patterns. If you haven't already, please verify that queries actually filter using all three columns in this order.openmeter/entitlement/balanceworker/ingesthandler.go (1)
21-25: LGTM! Nice simplification.The change from passing a slice of filters to a single filter makes the code cleaner and more straightforward. Since
MeterSlugsis still a slice, you maintain the ability to query multiple meters in one go, which is perfect for this use case.test/entitlement/regression/scenario_test.go (1)
634-640: LGTM!The test correctly adapts to the updated API signature, passing a single
IngestEventQueryFilterinstead of a slice. The assertions remain unchanged, which is exactly what we'd expect for this kind of interface refactor.openmeter/entitlement/balanceworker/repository.go (1)
12-29: LGTM!The interface simplification from accepting a slice of filters to a single filter is a nice cleanup. The streamlined
ListAffectedEntitlementsResponsestruct (removing redundant fields that can be derived elsewhere) makes the contract cleaner and easier to work with.openmeter/entitlement/adapter/entitlement.go (1)
297-309: LGTM!The response building logic is clean and correctly maps all the required fields. Good use of
convert.SafeToUTC()for the nullable time fields.
chrisgacsal
force-pushed
the
fix/listentitlement-whatever-by-whaterver-by-whatever
branch
from
443d3af to
d7c7f06
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI Agents
In @openmeter/entitlement/adapter/entitlement.go:
- Around line 285-298: The query constructs a Where clause using
repo.db.Entitlement.Query but unconditionally passes db_entitlement.CustomerIDIn
and db_entitlement.FeatureIDIn with slices that may be empty; change to build a
slice of predicates (e.g., preds) starting with
db_entitlement.Namespace(eventFilter.Namespace) and only append
db_entitlement.CustomerIDIn(lo.Uniq(...)) if len(customers) > 0 and only append
db_entitlement.FeatureIDIn(lo.Map(features,...)) if len(features) > 0, then call
.Where(preds...) on repo.db.Entitlement.Query before .All(ctx) to avoid passing
empty IN predicates.
📜 Review details
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (2)
openmeter/ent/db/migrate/schema.gois excluded by!**/ent/db/**tools/migrate/migrations/atlas.sumis excluded by!**/*.sum,!**/*.sum
📒 Files selected for processing (8)
openmeter/ent/schema/customer.goopenmeter/ent/schema/feature.goopenmeter/entitlement/adapter/entitlement.goopenmeter/entitlement/balanceworker/ingesthandler.goopenmeter/entitlement/balanceworker/repository.gotest/entitlement/regression/scenario_test.gotools/migrate/migrations/20260106184018_balance_worker_idx_improvements.down.sqltools/migrate/migrations/20260106184018_balance_worker_idx_improvements.up.sql
🚧 Files skipped from review as they are similar to previous changes (2)
- openmeter/ent/schema/customer.go
- openmeter/entitlement/balanceworker/repository.go
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go
⚙️ CodeRabbit configuration file
**/*.go: In general when reviewing the Golang code make readability and maintainability a priority, even potentially suggest restructuring the code to improve them.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.
Files:
openmeter/entitlement/adapter/entitlement.goopenmeter/entitlement/balanceworker/ingesthandler.gotest/entitlement/regression/scenario_test.goopenmeter/ent/schema/feature.go
**/*_test.go
⚙️ CodeRabbit configuration file
**/*_test.go: Make sure the tests are comprehensive and cover the changes. Keep a strong focus on unit tests and in-code integration tests.
When appropriate, recommend e2e tests for critical changes.
Files:
test/entitlement/regression/scenario_test.go
🧠 Learnings (2)
📚 Learning: 2025-04-20T11:15:07.499Z
Learnt from: chrisgacsal
Repo: openmeterio/openmeter PR: 2692
File: openmeter/productcatalog/plan/adapter/mapping.go:64-74
Timestamp: 2025-04-20T11:15:07.499Z
Learning: In the OpenMeter codebase, Ent's edge methods ending in "OrErr" (like AddonsOrErr()) only return NotLoadedError when the edge wasn't loaded, and cannot return DB errors. Simple err != nil checks are sufficient for these methods.
Applied to files:
openmeter/entitlement/adapter/entitlement.go
📚 Learning: 2025-10-09T13:59:12.012Z
Learnt from: chrisgacsal
Repo: openmeterio/openmeter PR: 3486
File: openmeter/ingest/kafkaingest/serializer/serializer.go:105-107
Timestamp: 2025-10-09T13:59:12.012Z
Learning: In OpenMeter, the CloudEvents `subject` field is mandatory for the application's business logic, even though it's optional in the CloudEvents specification. The `ValidateKafkaPayloadToCloudEvent` function in `openmeter/ingest/kafkaingest/serializer/serializer.go` intentionally enforces this requirement.
Applied to files:
test/entitlement/regression/scenario_test.go
🧬 Code graph analysis (4)
openmeter/entitlement/adapter/entitlement.go (2)
openmeter/entitlement/balanceworker/repository.go (2)
IngestEventQueryFilter(16-20)ListAffectedEntitlementsResponse(22-29)pkg/clock/clock.go (1)
Now(14-21)
openmeter/entitlement/balanceworker/ingesthandler.go (2)
openmeter/entitlement/balanceworker/repository.go (1)
IngestEventQueryFilter(16-20)openmeter/subject/subject.go (1)
SubjectKey(34-36)
test/entitlement/regression/scenario_test.go (2)
openmeter/entitlement/balanceworker/repository.go (1)
IngestEventQueryFilter(16-20)app/common/namespace.go (1)
Namespace(13-15)
openmeter/ent/schema/feature.go (1)
pkg/models/annotation.go (1)
Annotations(5-5)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
- GitHub Check: Artifacts / Benthos Collector Container image
- GitHub Check: Artifacts / Container image
- GitHub Check: Migration Checks
- GitHub Check: Build
- GitHub Check: Test
- GitHub Check: Lint
- GitHub Check: Code Generators
- GitHub Check: Repository Scan
- GitHub Check: Analyze (go)
🔇 Additional comments (7)
tools/migrate/migrations/20260106184018_balance_worker_idx_improvements.down.sql (1)
1-4: LGTM!Down migration correctly reverses the index creations from the up migration. The DROP INDEX statements are straightforward and will cleanly roll back the changes.
tools/migrate/migrations/20260106184018_balance_worker_idx_improvements.up.sql (1)
1-4: Indexes align well with the new query patterns.Both indexes support the optimized query flow in
ListEntitlementsAffectedByIngestEvents. The partial index on features witharchived_at IS NULLis a nice touch - keeps the index lean.One thing to keep in mind: these
CREATE INDEXstatements will acquire a lock on the tables during creation. For large production tables, you might want to considerCREATE INDEX CONCURRENTLYto avoid blocking writes. Just flagging in case the tables are sizable.openmeter/ent/schema/feature.go (1)
46-49: LGTM!The new index follows the same pattern as the existing
(namespace, key)index with thearchived_at IS NULLfilter. Non-unique is correct here since multiple features can reference the same meter slug. Nicely consistent with the codebase style.openmeter/entitlement/balanceworker/ingesthandler.go (1)
21-25: LGTM!Clean update to the new single-filter signature. The filter is properly constructed with all required fields from the ingest event. This simplifies the call site nicely.
test/entitlement/regression/scenario_test.go (1)
634-638: LGTM!Test correctly updated to match the new single-filter signature. Good that the expectations remain the same - this validates the refactored query still returns the expected entitlements.
openmeter/entitlement/adapter/entitlement.go (2)
252-257: Nice refactor to single-filter approach.The signature simplification makes sense - the call sites only ever pass a single filter. The consolidated query flow (resolve customers → resolve features → query entitlements) is cleaner than the previous iteration-based approach.
275-283: Error handling for features query is now in place.The
errcheck after the features query is properly added, preventing the silent error loss that was flagged in the previous review.
| entitlements, err := repo.db.Entitlement.Query(). | ||
| Where( | ||
| db_entitlement.Namespace(eventFilter.Namespace), | ||
| db_entitlement.CustomerIDIn(lo.Uniq(lo.Map(customers, func(item *db.Customer, _ int) string { | ||
| return item.ID | ||
| }))...), | ||
| db_entitlement.FeatureIDIn(lo.Map(features, func(item *db.Feature, _ int) string { | ||
| return item.ID | ||
| })...), | ||
| ). | ||
| All(ctx) | ||
| if err != nil { | ||
| return nil, fmt.Errorf("failed to query entitlements: %w", err) | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, let's find and read the actual code file
find . -name "entitlement.go" -path "*/entitlement/adapter/*" | head -5Repository: openmeterio/openmeter
Length of output: 111
🏁 Script executed:
# Search for the function mentioned in the review
rg -n "ListEntitlementsAffectedByIngestEvents" --type=go -B5 -A15 | head -150Repository: openmeterio/openmeter
Length of output: 7574
🏁 Script executed:
# Check how *In() predicates are typically used in the codebase
rg -n "CustomerIDIn\|FeatureIDIn" --type=go -B2 -A2 | head -80Repository: openmeterio/openmeter
Length of output: 47
🏁 Script executed:
# Read the entitlement.go file, focusing on the relevant section
sed -n '250,310p' openmeter/entitlement/adapter/entitlement.goRepository: openmeterio/openmeter
Length of output: 2496
🏁 Script executed:
# Also check if there are other similar queries in the codebase to see how empty slices are handled
rg "\.Query\(\).*Where.*\)" --type=go -A5 | grep -A5 "In(" | head -50Repository: openmeterio/openmeter
Length of output: 47
🏁 Script executed:
# Check the Ent codegen to understand how *In() predicates work with empty slices
find . -name "*.go" -path "*entgo*" -o -path "*ent*" | grep -i predicate | head -10Repository: openmeterio/openmeter
Length of output: 302
🏁 Script executed:
# Search for tests that might cover empty result scenarios
rg "ListEntitlementsAffectedByIngestEvents" --type=go -B10 -A30 | grep -E "(test|Test|assert|Len|empty|zero)" -iRepository: openmeterio/openmeter
Length of output: 3013
🏁 Script executed:
# Check for any defensive checks before calling *In() predicates
rg "CustomerIDIn|FeatureIDIn" --type=go -B5 -A5Repository: openmeterio/openmeter
Length of output: 3454
🏁 Script executed:
# Search for Ent predicate implementations to understand empty slice behavior
find . -name "predicate.go" -path "*entitlement*" | head -5Repository: openmeterio/openmeter
Length of output: 47
🏁 Script executed:
# Look for any test files that test the ListEntitlementsAffectedByIngestEvents function with edge cases
rg -l "ListEntitlementsAffectedByIngestEvents" --type=goRepository: openmeterio/openmeter
Length of output: 257
🏁 Script executed:
# Let me check if there are any issues or PRs mentioning Ent's handling of empty IN() predicates
rg "IN.*empty|empty.*IN" --type=go | head -20Repository: openmeterio/openmeter
Length of output: 47
🏁 Script executed:
# Let's look at more examples of defensive checks in the codebase for *In() predicates
rg "len\(" --type=go -A2 | grep -B2 "In(" | head -40Repository: openmeterio/openmeter
Length of output: 2254
🏁 Script executed:
# Search for any comments or documentation about why these defensive checks exist
rg -i "empty.*in|in.*empty" openmeter/subscription/repo/subscriptionrepo.go -B3 -A3Repository: openmeterio/openmeter
Length of output: 47
🏁 Script executed:
# Check if there are any related issues or TODOs in the entitlement file about this
rg -i "TODO|FIXME|BUG" openmeter/entitlement/adapter/entitlement.go | head -20Repository: openmeterio/openmeter
Length of output: 350
Add defensive checks before using *In() predicates with potentially empty slices.
The code calls CustomerIDIn() and FeatureIDIn() without checking if the slices are empty. Throughout the codebase, every other use of *In() predicates is guarded with if len(...) > 0 checks (see subscription/repo, productcatalog/planaddon/adapter, billing/adapter/invoice, etc.). Add the same defensive checks here to match the established pattern and avoid potential issues with empty IN predicates:
if len(customers) > 0 {
// Add CustomerIDIn to Where clause
}
if len(features) > 0 {
// Add FeatureIDIn to Where clause
}
🤖 Prompt for AI Agents
In @openmeter/entitlement/adapter/entitlement.go around lines 285 - 298, The
query constructs a Where clause using repo.db.Entitlement.Query but
unconditionally passes db_entitlement.CustomerIDIn and
db_entitlement.FeatureIDIn with slices that may be empty; change to build a
slice of predicates (e.g., preds) starting with
db_entitlement.Namespace(eventFilter.Namespace) and only append
db_entitlement.CustomerIDIn(lo.Uniq(...)) if len(customers) > 0 and only append
db_entitlement.FeatureIDIn(lo.Map(features,...)) if len(features) > 0, then call
.Where(preds...) on repo.db.Entitlement.Query before .All(ctx) to avoid passing
empty IN predicates.
|
😢 |
3 participants
Overview
Improve SQL query used for fetching entitlements by customers and corresponding meters.
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.