Merge pull request #1469 from openmeterio/feat/get-expired-entitlemen…

…ts-for-multiple-namespaces

feat: fetch expired entitlements for multiple namespaces

openmeter/ent/schema/entitlement.go

@@ -55,6 +55,8 @@ func (Entitlement) Indexes() []ent.Index {
 		index.Fields("namespace", "id", "subject_key"),
 		index.Fields("namespace", "feature_id", "id"),
 		index.Fields("namespace", "current_usage_period_end"),
+		// index for collecting all the entitlements with due resets
+		index.Fields("current_usage_period_end", "deleted_at"),
 	}
 }
 

openmeter/entitlement/adapter/entitlement.go

@@ -409,15 +409,19 @@ func (a *entitlementDBAdapter) UpdateEntitlementUsagePeriod(ctx context.Context,
 	return err
 }
 
-func (a *entitlementDBAdapter) ListEntitlementsWithExpiredUsagePeriod(ctx context.Context, namespace string, expiredBefore time.Time) ([]entitlement.Entitlement, error) {
-	res, err := withLatestUsageReset(a.db.Entitlement.Query(), []string{namespace}).
+func (a *entitlementDBAdapter) ListEntitlementsWithExpiredUsagePeriod(ctx context.Context, namespaces []string, expiredBefore time.Time) ([]entitlement.Entitlement, error) {
+	query := withLatestUsageReset(a.db.Entitlement.Query(), namespaces).
 		Where(
-			db_entitlement.Namespace(namespace),
 			db_entitlement.CurrentUsagePeriodEndNotNil(),
 			db_entitlement.CurrentUsagePeriodEndLTE(expiredBefore),
 			db_entitlement.Or(db_entitlement.DeletedAtIsNil(), db_entitlement.DeletedAtGT(clock.Now())),
-		).
-		All(ctx)
+		)
+
+	if len(namespaces) > 0 {
+		query = query.Where(db_entitlement.NamespaceIn(namespaces...))
+	}
+
+	res, err := query.All(ctx)
 	if err != nil {
 		return nil, err
 	}

openmeter/entitlement/metered/reset.go

@@ -73,7 +73,7 @@ func (e *connector) ResetEntitlementUsage(ctx context.Context, entitlementID mod
 }
 
 func (c *connector) ResetEntitlementsWithExpiredUsagePeriod(ctx context.Context, namespace string, highwatermark time.Time) ([]models.NamespacedID, error) {
-	entitlements, err := c.entitlementRepo.ListEntitlementsWithExpiredUsagePeriod(ctx, namespace, highwatermark)
+	entitlements, err := c.entitlementRepo.ListEntitlementsWithExpiredUsagePeriod(ctx, []string{namespace}, highwatermark)
 	if err != nil {
 		return nil, fmt.Errorf("failed to list entitlements with due reset: %w", err)
 	}

openmeter/entitlement/repository.go

@@ -32,7 +32,7 @@ type EntitlementRepo interface {
 	LockEntitlementForTx(ctx context.Context, entitlementID models.NamespacedID) error
 
 	UpdateEntitlementUsagePeriod(ctx context.Context, entitlementID models.NamespacedID, params UpdateEntitlementUsagePeriodParams) error
-	ListEntitlementsWithExpiredUsagePeriod(ctx context.Context, namespace string, highwatermark time.Time) ([]Entitlement, error)
+	ListEntitlementsWithExpiredUsagePeriod(ctx context.Context, namespaces []string, highwatermark time.Time) ([]Entitlement, error)
 
 	entutils.TxCreator
 	entutils.TxUser[EntitlementRepo]

tools/migrate/migrations/20240903155435_entitlement-expired-index.down.sql

@@ -0,0 +1,2 @@
+-- reverse: create index "entitlement_current_usage_period_end_deleted_at" to table: "entitlements"
+DROP INDEX "entitlement_current_usage_period_end_deleted_at";

tools/migrate/migrations/20240903155435_entitlement-expired-index.up.sql

@@ -0,0 +1,2 @@
+-- create index "entitlement_current_usage_period_end_deleted_at" to table: "entitlements"
+CREATE INDEX "entitlement_current_usage_period_end_deleted_at" ON "entitlements" ("current_usage_period_end", "deleted_at");

tools/migrate/migrations/atlas.sum

@@ -1,3 +1,5 @@
-h1:EPap+DaNoXsLQC61DQvaBbCLBcYaGP1DHSglnzexChY=
+h1:jmrq8Fzq+/Ds0gmHmWUNruJ3jzKFhVVBYSDlbtordSA=
 20240826120919_init.down.sql h1:AIbgwwngjkJEYa3yRZsIXQyBa2+qoZttwMXHxXEbHLI=
 20240826120919_init.up.sql h1:/hYHWF3Z3dab8SMKnw99ixVktCuJe2bAw5wstCZIEN8=
+20240903155435_entitlement-expired-index.down.sql h1:np2xgYs3KQ2z7qPBcobtGNhqWQ3V8NwEP9E5U3TmpSA=
+20240903155435_entitlement-expired-index.up.sql h1:HegNgovI3MQxQMd/072kOchfm5B37SXzODZIteOwX78=