feat(customer): get by usage attribution lookup by customer key and id

Author: hekike

openmeter/customer/adapter/customer.go

@@ -392,7 +392,15 @@ func (a *adapter) GetCustomerByUsageAttribution(ctx context.Context, input custo
 
 	query := a.db.Customer.Query().
 		Where(customerdb.Namespace(input.Namespace)).
-		Where(customerdb.HasSubjectsWith(customersubjectsdb.SubjectKey(input.SubjectKey))).
+		Where(
+			customerdb.Or(
+				// We lookup the customer by subject key in the subjects table
+				customerdb.HasSubjectsWith(customersubjectsdb.SubjectKey(input.SubjectKey)),
+				// Or we lookup the customer by ID or key
+				customerdb.ID(input.SubjectKey),
+				customerdb.Key(input.SubjectKey),
+			),
+		).
 		Where(customerdb.DeletedAtIsNil())
 	query = withSubjects(query)
 	query = withActiveSubscription(query)

test/customer/customer.go

@@ -623,6 +623,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,
 			},
@@ -642,8 +643,28 @@ func (s *CustomerHandlerTestSuite) TestGetByUsageAttribution(ctx context.Context
 	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 key
+	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
+		Namespace:  s.namespace,
+		SubjectKey: 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,
+		SubjectKey: 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 usage attribution with a non-existent subject key
 	_, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{