chore: improve entitlement lookup performance (#4149)

Author: turip

openmeter/entitlement/adapter/entitlement.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"entgo.io/ent/dialect/sql"
+	"github.com/oklog/ulid/v2"
 	"github.com/samber/lo"
 
 	customeradapter "github.com/openmeterio/openmeter/openmeter/customer/adapter"
@@ -86,10 +87,9 @@ func (a *entitlementDBAdapter) GetActiveEntitlementOfCustomerAt(ctx context.Cont
 				Where(EntitlementActiveAt(at)...).
 				Where(
 					db_entitlement.Or(db_entitlement.DeletedAtGT(at), db_entitlement.DeletedAtIsNil()),
+					db_entitlement.CustomerID(customerID),
 					db_entitlement.HasCustomerWith(
-						customerdb.Namespace(namespace),
 						customerNotDeletedAt(at),
-						customerdb.ID(customerID),
 					),
 					db_entitlement.Namespace(namespace),
 					db_entitlement.FeatureKey(featureKey),
@@ -313,6 +313,9 @@ func (a *entitlementDBAdapter) ListEntitlements(ctx context.Context, params enti
 		a,
 		func(ctx context.Context, repo *entitlementDBAdapter) (pagination.Result[entitlement.Entitlement], error) {
 			now := clock.Now().UTC()
+			response := pagination.Result[entitlement.Entitlement]{
+				Page: params.Page,
+			}
 
 			query := repo.db.Entitlement.Query()
 
@@ -335,15 +338,29 @@ func (a *entitlementDBAdapter) ListEntitlements(ctx context.Context, params enti
 			}
 
 			if len(params.CustomerKeys) > 0 {
-				query = query.Where(db_entitlement.HasCustomerWith(
-					customerdb.KeyIn(params.CustomerKeys...),
-				))
+				customerQuery := repo.db.Customer.Query().
+					Where(customerdb.KeyIn(params.CustomerKeys...))
+
+				if len(params.Namespaces) > 0 {
+					customerQuery = customerQuery.Where(customerdb.NamespaceIn(params.Namespaces...))
+				}
+
+				customerIDs, err := customerQuery.IDs(ctx)
+				if err != nil {
+					return response, err
+				}
+
+				if len(customerIDs) == 0 {
+					response.Items = []entitlement.Entitlement{}
+					response.TotalCount = 0
+					return response, nil
+				}
+
+				query = query.Where(db_entitlement.CustomerIDIn(customerIDs...))
 			}
 
 			if len(params.CustomerIDs) > 0 {
-				query = query.Where(db_entitlement.HasCustomerWith(
-					customerdb.IDIn(params.CustomerIDs...),
-				))
+				query = query.Where(db_entitlement.CustomerIDIn(params.CustomerIDs...))
 			}
 
 			if len(params.EntitlementTypes) > 0 {
@@ -357,16 +374,30 @@ func (a *entitlementDBAdapter) ListEntitlements(ctx context.Context, params enti
 			}
 
 			if len(params.FeatureIDsOrKeys) > 0 {
-				var ep predicate.Entitlement
-				for i, idOrKey := range params.FeatureIDsOrKeys {
-					p := db_entitlement.Or(db_entitlement.FeatureID(idOrKey), db_entitlement.FeatureKey(idOrKey))
-					if i == 0 {
-						ep = p
-						continue
+				featureIDs := make([]string, 0, len(params.FeatureIDsOrKeys))
+				featureKeys := make([]string, 0, len(params.FeatureIDsOrKeys))
+
+				for _, idOrKey := range params.FeatureIDsOrKeys {
+					if _, err := ulid.Parse(idOrKey); err == nil {
+						featureIDs = append(featureIDs, idOrKey)
+					} else {
+						featureKeys = append(featureKeys, idOrKey)
 					}
-					ep = db_entitlement.Or(ep, p)
 				}
-				query = query.Where(ep)
+
+				switch {
+				case len(featureIDs) > 0 && len(featureKeys) > 0:
+					query = query.Where(
+						db_entitlement.Or(
+							db_entitlement.FeatureIDIn(featureIDs...),
+							db_entitlement.FeatureKeyIn(featureKeys...),
+						),
+					)
+				case len(featureIDs) > 0:
+					query = query.Where(db_entitlement.FeatureIDIn(featureIDs...))
+				case len(featureKeys) > 0:
+					query = query.Where(db_entitlement.FeatureKeyIn(featureKeys...))
+				}
 			}
 
 			if len(params.FeatureIDs) > 0 {
@@ -406,10 +437,6 @@ func (a *entitlementDBAdapter) ListEntitlements(ctx context.Context, params enti
 				}
 			}
 
-			response := pagination.Result[entitlement.Entitlement]{
-				Page: params.Page,
-			}
-
 			// we're using limit and offset
 			if params.Page.IsZero() {
 				if params.Limit > 0 {

openmeter/entitlement/adapter/entitlement_test.go

@@ -117,6 +117,33 @@ func createCustomerWithSubject(t *testing.T, subjectRepo subject.Service, custom
 	return cust
 }
 
+func createCustomerWithSubjectAndKey(t *testing.T, subjectRepo subject.Service, customerRepo customer.Adapter, namespace string, subjectKey string, customerKey string) *customer.Customer {
+	t.Helper()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	_, err := subjectRepo.Create(ctx, subject.CreateInput{
+		Namespace: namespace,
+		Key:       subjectKey,
+	})
+	require.NoError(t, err)
+
+	cust, err := customerRepo.CreateCustomer(ctx, customer.CreateCustomerInput{
+		Namespace: namespace,
+		CustomerMutate: customer.CustomerMutate{
+			Key:  lo.ToPtr(customerKey),
+			Name: "Customer 1",
+			UsageAttribution: &customer.CustomerUsageAttribution{
+				SubjectKeys: []string{subjectKey},
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	return cust
+}
+
 func TestUpsertEntitlementCurrentPeriods(t *testing.T) {
 	ns := "ns1"
 	featureKey := "feature1"
@@ -518,3 +545,78 @@ func TestEntitlementLoadsSubjectAndCustomerAndPreservesAcrossTypedMapping(t *tes
 		}
 	}
 }
+
+func TestListEntitlementsFiltersByCustomerKeysAndFeatureIDsOrKeys(t *testing.T) {
+	ctx := context.Background()
+	ns := "ns-list-filters"
+
+	repo, cleanup := setup(t)
+	defer cleanup()
+
+	featureByKey, err := repo.featureRepo.CreateFeature(ctx, feature.CreateFeatureInputs{
+		Namespace: ns,
+		Key:       "free_plan_usage",
+		Name:      "Free plan usage",
+	})
+	require.NoError(t, err)
+
+	featureByID, err := repo.featureRepo.CreateFeature(ctx, feature.CreateFeatureInputs{
+		Namespace: ns,
+		Key:       "pro_plan_usage",
+		Name:      "Pro plan usage",
+	})
+	require.NoError(t, err)
+
+	customerA := createCustomerWithSubjectAndKey(t, repo.subjectRepo, repo.customerRepo, ns, "subject-a", "customer-a")
+	customerB := createCustomerWithSubjectAndKey(t, repo.subjectRepo, repo.customerRepo, ns, "subject-b", "customer-b")
+
+	entA, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
+		Namespace:        ns,
+		FeatureID:        featureByKey.ID,
+		FeatureKey:       featureByKey.Key,
+		UsageAttribution: customerA.GetUsageAttribution(),
+		EntitlementType:  entitlement.EntitlementTypeBoolean,
+	})
+	require.NoError(t, err)
+
+	entB, err := repo.entRepo.CreateEntitlement(ctx, entitlement.CreateEntitlementRepoInputs{
+		Namespace:        ns,
+		FeatureID:        featureByID.ID,
+		FeatureKey:       featureByID.Key,
+		UsageAttribution: customerB.GetUsageAttribution(),
+		EntitlementType:  entitlement.EntitlementTypeBoolean,
+	})
+	require.NoError(t, err)
+
+	t.Run("Should filter by customer key and feature key", func(t *testing.T) {
+		res, err := repo.entRepo.ListEntitlements(ctx, entitlement.ListEntitlementsParams{
+			Namespaces:       []string{ns},
+			CustomerKeys:     []string{"customer-a"},
+			FeatureIDsOrKeys: []string{featureByKey.Key},
+		})
+		require.NoError(t, err)
+		require.Len(t, res.Items, 1)
+		require.Equal(t, entA.ID, res.Items[0].ID)
+	})
+
+	t.Run("Should filter by customer key and feature ID", func(t *testing.T) {
+		res, err := repo.entRepo.ListEntitlements(ctx, entitlement.ListEntitlementsParams{
+			Namespaces:       []string{ns},
+			CustomerKeys:     []string{"customer-b"},
+			FeatureIDsOrKeys: []string{featureByID.ID},
+		})
+		require.NoError(t, err)
+		require.Len(t, res.Items, 1)
+		require.Equal(t, entB.ID, res.Items[0].ID)
+	})
+
+	t.Run("Should return empty result without querying entitlements when customer key does not exist", func(t *testing.T) {
+		res, err := repo.entRepo.ListEntitlements(ctx, entitlement.ListEntitlementsParams{
+			Namespaces:   []string{ns},
+			CustomerKeys: []string{"missing-customer"},
+		})
+		require.NoError(t, err)
+		require.Empty(t, res.Items)
+		require.Zero(t, res.TotalCount)
+	})
+}

openmeter/entitlement/service/service.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/oklog/ulid/v2"
 	"github.com/samber/lo"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/sync/semaphore"
@@ -226,10 +227,17 @@ func (c *service) GetEntitlementsOfCustomer(ctx context.Context, namespace strin
 }
 
 func (c *service) GetEntitlementOfCustomerAt(ctx context.Context, namespace string, customerID string, idOrFeatureKey string, at time.Time) (*entitlement.Entitlement, error) {
+	// Feature keys are forbidden from being valid ULIDs, so non-ULID inputs can
+	// skip the guaranteed-miss entitlement-ID lookup.
+	if _, err := ulid.Parse(idOrFeatureKey); err != nil {
+		return c.entitlementRepo.GetActiveEntitlementOfCustomerAt(ctx, namespace, customerID, idOrFeatureKey, at)
+	}
+
 	ent, err := c.entitlementRepo.GetEntitlement(ctx, models.NamespacedID{Namespace: namespace, ID: idOrFeatureKey})
 	if _, ok := lo.ErrorsAs[*entitlement.NotFoundError](err); ok {
 		ent, err = c.entitlementRepo.GetActiveEntitlementOfCustomerAt(ctx, namespace, customerID, idOrFeatureKey, at)
 	}
+
 	return ent, err
 }
 

openmeter/entitlement/service/service_test.go

@@ -0,0 +1,80 @@
+package service_test
+
+import (
+	"testing"
+	"time"
+
+	"github.com/samber/lo"
+	"github.com/stretchr/testify/require"
+
+	"github.com/openmeterio/openmeter/openmeter/entitlement"
+	"github.com/openmeterio/openmeter/openmeter/meter"
+	"github.com/openmeterio/openmeter/openmeter/productcatalog/feature"
+	"github.com/openmeterio/openmeter/openmeter/testutils"
+	"github.com/openmeterio/openmeter/pkg/clock"
+)
+
+func TestGetEntitlementOfCustomerAt(t *testing.T) {
+	conn, deps := setupDependecies(t)
+	defer deps.Teardown()
+
+	namespace := "ns-get-entitlement-of-customer-at"
+	now := testutils.GetRFC3339Time(t, "2025-01-01T00:00:00Z")
+
+	clock.SetTime(now)
+	defer clock.ResetTime()
+
+	mtr, err := deps.meterService.CreateMeter(t.Context(), meter.CreateMeterInput{
+		Namespace:     namespace,
+		Name:          "Meter 1",
+		Key:           "meter1",
+		Description:   nil,
+		Aggregation:   meter.MeterAggregationSum,
+		EventType:     "test",
+		EventFrom:     nil,
+		ValueProperty: lo.ToPtr("$.value"),
+		GroupBy:       nil,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, mtr)
+	createMeterInPG(t, deps.dbClient, mtr)
+
+	cust := createCustomerAndSubject(t, deps.subjectService, deps.customerService, namespace, "cust-1", "Customer 1")
+
+	feat, err := deps.featureRepo.CreateFeature(t.Context(), feature.CreateFeatureInputs{
+		Key:       "free_plan_usage",
+		Name:      "Free plan usage",
+		Namespace: namespace,
+		MeterID:   &mtr.ID,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, feat)
+
+	ent, err := conn.CreateEntitlement(t.Context(), entitlement.CreateEntitlementInputs{
+		Namespace:        namespace,
+		UsageAttribution: cust.GetUsageAttribution(),
+		FeatureKey:       &feat.Key,
+		FeatureID:        &feat.ID,
+		EntitlementType:  entitlement.EntitlementTypeBoolean,
+	}, nil)
+	require.NoError(t, err)
+	require.NotNil(t, ent)
+
+	t.Run("Should resolve entitlement by feature key", func(t *testing.T) {
+		res, err := conn.GetEntitlementOfCustomerAt(t.Context(), namespace, cust.ID, feat.Key, clock.Now().Add(time.Hour))
+		require.NoError(t, err)
+		require.NotNil(t, res)
+		require.Equal(t, ent.ID, res.ID)
+		require.Equal(t, feat.Key, res.FeatureKey)
+		require.Equal(t, cust.ID, res.CustomerID)
+	})
+
+	t.Run("Should resolve entitlement by entitlement ID", func(t *testing.T) {
+		res, err := conn.GetEntitlementOfCustomerAt(t.Context(), namespace, cust.ID, ent.ID, clock.Now().Add(time.Hour))
+		require.NoError(t, err)
+		require.NotNil(t, res)
+		require.Equal(t, ent.ID, res.ID)
+		require.Equal(t, feat.Key, res.FeatureKey)
+		require.Equal(t, cust.ID, res.CustomerID)
+	})
+}