feat(streaming): lookup usage by customer key and subjects

openmeter/customer/adapter/customer.go

@@ -462,13 +462,20 @@ func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input custo
 
 		query := repo.db.Customer.Query().
 			Where(customerdb.Namespace(input.Namespace)).
-			Where(customerdb.HasSubjectsWith(
-				customersubjectsdb.SubjectKey(input.SubjectKey),
-				customersubjectsdb.Or(
-					customersubjectsdb.DeletedAtIsNil(),
-					customersubjectsdb.DeletedAtGT(now),
+			Where(
+				customerdb.Or(
+					// We lookup the customer by subject key in the subjects table
+					customerdb.HasSubjectsWith(
+						customersubjectsdb.SubjectKey(input.Key),
+						customersubjectsdb.Or(
+							customersubjectsdb.DeletedAtIsNil(),
+							customersubjectsdb.DeletedAtGT(now),
+						),
+					),
+					// Or else we lookup the customer by key in the customers table
+					customerdb.Key(input.Key),
 				),
-			)).
+			).
 			Where(customerdb.DeletedAtIsNil())
 		query = WithSubjects(query, now)
 		if slices.Contains(input.Expands, customer.ExpandSubscriptions) {
@@ -479,7 +486,7 @@ func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input custo
 		if err != nil {
 			if entdb.IsNotFound(err) {
 				return nil, models.NewGenericNotFoundError(
-					fmt.Errorf("customer with subject key %s not found in %s namespace", input.SubjectKey, input.Namespace),
+					fmt.Errorf("customer with subject key %s not found in %s namespace", input.Key, input.Namespace),
 				)
 			}
 

openmeter/customer/customer.go

@@ -231,8 +231,10 @@ func (c CustomerUsageAttribution) GetFirstSubjectKey() (string, error) {
 
 // GetCustomerByUsageAttributionInput represents the input for the GetCustomerByUsageAttribution method
 type GetCustomerByUsageAttributionInput struct {
-	Namespace  string
-	SubjectKey string
+	Namespace string
+
+	// The key of either the customer or one of its subjects
+	Key string
 
 	// Expand
 	Expands Expands
@@ -243,7 +245,7 @@ func (i GetCustomerByUsageAttributionInput) Validate() error {
 		return models.NewGenericValidationError(errors.New("namespace is required"))
 	}
 
-	if i.SubjectKey == "" {
+	if i.Key == "" {
 		return models.NewGenericValidationError(errors.New("subject key is required"))
 	}
 

openmeter/customer/service/hooks/subjectcustomer.go

@@ -68,8 +68,8 @@ func (s subjectCustomerHook) PostDelete(ctx context.Context, sub *subject.Subjec
 
 	// Let's get the customer by usage attribution
 	cus, err := s.provisioner.customer.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  sub.Namespace,
-		SubjectKey: sub.Key,
+		Namespace: sub.Namespace,
+		Key:       sub.Key,
 	})
 	if err != nil {
 		if models.IsGenericNotFoundError(err) {
@@ -305,8 +305,8 @@ var ErrCustomerKeyConflict = errors.New("customer key conflict")
 func (p CustomerProvisioner) getCustomerForSubject(ctx context.Context, sub *subject.Subject) (*customer.Customer, error) {
 	// Try to find Customer for Subject by usage attribution
 	cus, err := p.customer.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  sub.Namespace,
-		SubjectKey: sub.Key,
+		Namespace: sub.Namespace,
+		Key:       sub.Key,
 	})
 	if err != nil && !models.IsGenericNotFoundError(err) {
 		return nil, err

openmeter/customer/service/hooks/subjectcustomer_test.go

@@ -150,7 +150,11 @@ func TestCustomerProvisioner_EnsureCustomer(t *testing.T) {
 			require.NoError(t, err, "creating subject should not fail")
 			assert.NotNilf(t, sub, "subject must not be nil")
 
-			cus, err := env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{
+			cusForSubject, err := provisioner.EnsureCustomer(ctx, &sub)
+			require.NoError(t, err, "provisioning customer should not fail")
+			assert.NotNilf(t, cusForSubject, "customer must not be nil")
+
+			_, err = env.CustomerService.CreateCustomer(ctx, customer.CreateCustomerInput{
 				Namespace: namespace,
 				CustomerMutate: customer.CustomerMutate{
 					Key:         lo.ToPtr(sub.Key),
@@ -170,22 +174,8 @@ func TestCustomerProvisioner_EnsureCustomer(t *testing.T) {
 					Annotation: nil,
 				},
 			})
-			require.NoError(t, err, "creating customer should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
-
-			cus, err = provisioner.EnsureCustomer(ctx, &sub)
-			require.NoError(t, err, "provisioning customer should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
 
-			cus, err = env.CustomerService.GetCustomer(ctx, customer.GetCustomerInput{
-				CustomerID: &customer.CustomerID{
-					Namespace: cus.Namespace,
-					ID:        cus.ID,
-				},
-			})
-			require.NoErrorf(t, err, "getting customer for subject should not fail")
-			assert.NotNilf(t, cus, "customer must not be nil")
-			AssertSubjectCustomerEqual(t, &sub, cus)
+			require.True(t, models.IsGenericConflictError(err), "creating customer should fail with conflict")
 		})
 
 		t.Run("CustomerKeyMismatch", func(t *testing.T) {

openmeter/customer/service/service_test.go

@@ -86,8 +86,8 @@ func Test_CustomerService(t *testing.T) {
 
 				t.Run("ByUsageAttribution", func(t *testing.T) {
 					cusByUsage, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-						Namespace:  cus.Namespace,
-						SubjectKey: cus.UsageAttribution.SubjectKeys[0],
+						Namespace: cus.Namespace,
+						Key:       cus.UsageAttribution.SubjectKeys[0],
 					})
 					require.NoError(t, err, "getting customer usage attribution should not fail")
 					assert.NotNilf(t, cusByUsage, "customer must not be nil")
@@ -128,8 +128,8 @@ func Test_CustomerService(t *testing.T) {
 
 				t.Run("ByUsageAttribution", func(t *testing.T) {
 					cusByUsage, err := env.CustomerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-						Namespace:  cus.Namespace,
-						SubjectKey: subjectKeys[1],
+						Namespace: cus.Namespace,
+						Key:       subjectKeys[1],
 					})
 					require.NoError(t, err, "getting customer usage attribution should not fail")
 					assert.NotNilf(t, cusByUsage, "customer must not be nil")

openmeter/entitlement/driver/entitlement.go

@@ -569,8 +569,8 @@ func (h *entitlementHandler) resolveCustomerFromSubject(ctx context.Context, nam
 	}
 
 	cust, err := h.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  namespace,
-		SubjectKey: subj.Key,
+		Namespace: namespace,
+		Key:       subj.Key,
 	})
 	if err != nil {
 		return nil, err

openmeter/entitlement/driver/metered.go

@@ -395,8 +395,8 @@ func (h *meteredEntitlementHandler) resolveCustomerFromSubject(ctx context.Conte
 	}
 
 	cust, err := h.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  namespace,
-		SubjectKey: subj.Key,
+		Namespace: namespace,
+		Key:       subj.Key,
 	})
 	if err != nil {
 		return nil, err

openmeter/meterevent/adapter/event.go

@@ -277,8 +277,8 @@ func (a *adapter) enrichEventsWithCustomerID(ctx context.Context, namespace stri
 		// FIXME: do this in a batches to avoid hitting the database for each event
 		// Get the customer by usage attribution subject key
 		cust, err := a.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-			Namespace:  namespace,
-			SubjectKey: event.Subject,
+			Namespace: namespace,
+			Key:       event.Subject,
 		})
 		if err != nil {
 			if models.IsGenericNotFoundError(err) {

openmeter/streaming/clickhouse/meter_query_test.go

@@ -375,6 +375,7 @@ func TestQueryMeter(t *testing.T) {
 							},
 							ID: "customer1",
 						},
+						Key: lo.ToPtr("customer-key-1"),
 						UsageAttribution: customer.CustomerUsageAttribution{
 							SubjectKeys: []string{"subject1"},
 						},
@@ -392,10 +393,17 @@ func TestQueryMeter(t *testing.T) {
 					},
 				},
 			},
-			wantSQL:  "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND om_events.subject IN (?)",
-			wantArgs: []interface{}{"my_namespace", "event1", []string{"subject1", "subject2"}},
+			wantSQL: "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND om_events.subject IN (?)",
+			wantArgs: []interface{}{"my_namespace", "event1", []string{
+				// Only the first customer has a key
+				"customer-key-1",
+				// Usage attribution subjects of the first customer
+				"subject1",
+				// Usage attribution subjects of the second customer
+				"subject2",
+			}},
 		},
-		{
+		{ // Filter by both customer and subject
 			name: "Filter by both customer and subject",
 			query: queryMeter{
 				Database:        "openmeter",

openmeter/streaming/clickhouse/queryhelper.go

@@ -5,11 +5,12 @@ import (
 	"strings"
 
 	"github.com/huandu/go-sqlbuilder"
-	"github.com/samber/lo"
 
 	"github.com/openmeterio/openmeter/openmeter/streaming"
 )
 
+const subjectToCustomerIDDictionary = "subject_to_customer_id"
+
 // selectCustomerIdColumn
 func selectCustomerIdColumn(eventsTableName string, customers []streaming.Customer, query *sqlbuilder.SelectBuilder) *sqlbuilder.SelectBuilder {
 	// If there are no customers, we return an empty customer id column
@@ -21,28 +22,43 @@ func selectCustomerIdColumn(eventsTableName string, customers []streaming.Custom
 	getColumn := columnFactory(eventsTableName)
 	subjectColumn := getColumn("subject")
 
-	// Build a map of subject to customer id
+	// Build a map of event subjects to customer ids
 	var values []string
 
+	// For each customer, we map event subjects to customer ids
 	for _, customer := range customers {
-		// Add each subject key to the map and map it to the customer id
+		customerIDSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(customer.GetUsageAttribution().ID))
+
+		// We map the customer key to the customer id if it exists
+		if customer.GetUsageAttribution().Key != nil {
+			customerKeySQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(*customer.GetUsageAttribution().Key))
+			values = append(values, customerKeySQL, customerIDSQL)
+		}
+
+		// We map each subject key to the customer id
 		for _, subjectKey := range customer.GetUsageAttribution().SubjectKeys {
 			subjectSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(subjectKey))
-			customerIDSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(customer.GetUsageAttribution().ID))
 
 			values = append(values, subjectSQL, customerIDSQL)
 		}
 	}
 
-	mapAs := "subject_to_customer_id"
-	mapSQL := fmt.Sprintf("WITH map(%s) as %s", strings.Join(values, ", "), mapAs)
+	// If there are no values, we return an empty customer id column
+	// This can happen if none of the customers has key or usage attribution subjects
+	if len(values) == 0 {
+		return query.SelectMore("'' AS customer_id")
+	}
+
+	// Name of the map (dictionary)
+
+	mapSQL := fmt.Sprintf("WITH map(%s) as %s", strings.Join(values, ", "), subjectToCustomerIDDictionary)
 
 	// Add the map to query via WITH clause
 	mapQuery := sqlbuilder.ClickHouse.NewCTEBuilder().SQL(mapSQL)
 	query = query.With(mapQuery)
 
 	// Select the customer id column
-	query = query.SelectMore(fmt.Sprintf("%s[%s] AS customer_id", mapAs, subjectColumn))
+	query = query.SelectMore(fmt.Sprintf("%s[%s] AS customer_id", subjectToCustomerIDDictionary, subjectColumn))
 
 	return query
 }
@@ -58,12 +74,26 @@ func customersWhere(eventsTableName string, customers []streaming.Customer, quer
 	getColumn := columnFactory(eventsTableName)
 	subjectColumn := getColumn("subject")
 
-	// If the customer filter is provided, we add all the subjects to the filter
-	subjects := lo.Map(customers, func(customer streaming.Customer, _ int) []string {
-		return customer.GetUsageAttribution().SubjectKeys
-	})
+	var subjects []string
+
+	// Collect all the subjects from the customers
+	for _, customer := range customers {
+		// Add the customer key to the filter if it exists
+		if customer.GetUsageAttribution().Key != nil {
+			subjects = append(subjects, *customer.GetUsageAttribution().Key)
+		}
+
+		// Add each subject key to the filter
+		subjects = append(subjects, customer.GetUsageAttribution().SubjectKeys...)
+	}
+
+	// If there are no subjects, we return an empty subject filter
+	// This can happen if none of the customers has key or usage attribution subjects
+	if len(subjects) == 0 {
+		return query
+	}
 
-	return query.Where(query.In(subjectColumn, lo.Flatten(subjects)))
+	return query.Where(query.In(subjectColumn, subjects))
 }
 
 // subjectWhere applies the subject filter to the query.

test/customer/customer.go

@@ -753,6 +753,7 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 		Namespace: s.namespace,
 		CustomerMutate: customer.CustomerMutate{
 			Name: TestName,
+			Key:  lo.ToPtr(TestKey),
 			UsageAttribution: customer.CustomerUsageAttribution{
 				SubjectKeys: TestSubjectKeys,
 			},
@@ -764,21 +765,41 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 
 	// Get the customer by usage attribution
 	cus, err := service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: TestSubjectKeys[0],
+		Namespace: s.namespace,
+		Key:       TestSubjectKeys[0],
+	})
+
+	require.NoError(t, err, "Fetching customer must not return error")
+	require.NotNil(t, cus, "Customer must not be nil")
+	require.Equal(t, s.namespace, cus.Namespace, "Customer namespace must match")
+	require.Equal(t, createdCustomer.ID, cus.ID, "Customer ID must match")
+
+	// Get the customer by key
+	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
+		Namespace: s.namespace,
+		Key:       TestKey,
+	})
+
+	require.NoError(t, err, "Fetching customer must not return error")
+	require.NotNil(t, cus, "Customer must not be nil")
+	require.Equal(t, s.namespace, cus.Namespace, "Customer namespace must match")
+	require.Equal(t, createdCustomer.ID, cus.ID, "Customer ID must match")
+
+	// Get the customer by key
+	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
+		Namespace: s.namespace,
+		Key:       TestKey,
 	})
 
 	require.NoError(t, err, "Fetching customer must not return error")
 	require.NotNil(t, cus, "Customer must not be nil")
 	require.Equal(t, s.namespace, cus.Namespace, "Customer namespace must match")
 	require.Equal(t, createdCustomer.ID, cus.ID, "Customer ID must match")
-	require.Equal(t, TestName, cus.Name, "Customer name must match")
-	require.Equal(t, TestSubjectKeys, cus.UsageAttribution.SubjectKeys, "Customer usage attribution subject keys must match")
 
 	// Get the customer by usage attribution with a non-existent subject key
 	_, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: "non-existent-subject-key",
+		Namespace: s.namespace,
+		Key:       "non-existent-subject-key",
 	})
 
 	require.True(t, models.IsGenericNotFoundError(err), "Fetching customer with non-existent subject key must return not found error")