refactor(customer): return map[string]*Customer from GetCustomersByUsageAttribution

Align the bulk resolver's semantics with the single-key lookup and the repo's
featureresolver.BatchResolve convention: every input key is present in the
returned map, with a nil value marking not-found (was: value map, absent =
not-found). Callers can distinguish requested-but-missing from not-requested,
and the map holds pointers instead of copied Customer structs.

Author: gergely-kurucz-konghq

openmeter/customer/service.go

@@ -26,6 +26,6 @@ type CustomerService interface {
 	DeleteCustomer(ctx context.Context, customer DeleteCustomerInput) error
 	GetCustomer(ctx context.Context, customer GetCustomerInput) (*Customer, error)
 	GetCustomerByUsageAttribution(ctx context.Context, input GetCustomerByUsageAttributionInput) (*Customer, error)
-	GetCustomersByUsageAttribution(ctx context.Context, input GetCustomersByUsageAttributionInput) (map[string]Customer, error)
+	GetCustomersByUsageAttribution(ctx context.Context, input GetCustomersByUsageAttributionInput) (map[string]*Customer, error)
 	UpdateCustomer(ctx context.Context, params UpdateCustomerInput) (*Customer, error)
 }

openmeter/customer/service/customer.go

@@ -150,19 +150,20 @@ func (s *Service) GetCustomerByUsageAttribution(ctx context.Context, input custo
 		return nil, err
 	}
 
-	c, ok := resolved[input.Key]
-	if !ok {
+	c := resolved[input.Key]
+	if c == nil {
 		return nil, models.NewGenericNotFoundError(
 			fmt.Errorf("customer with subject key %s not found in %s namespace", input.Key, input.Namespace),
 		)
 	}
 
-	return &c, nil
+	return c, nil
 }
 
 // GetCustomersByUsageAttribution resolves multiple customers by usage attribution keys in a single
 // query, mapping each input key to the customer it matches with key-over-subject precedence applied.
-func (s *Service) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]customer.Customer, error) {
+// Every input key is present in the returned map; keys that match no customer map to a nil value.
+func (s *Service) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) {
 	if err := input.Validate(); err != nil {
 		return nil, models.NewGenericValidationError(
 			fmt.Errorf("error getting customers by usage attribution: %w", err),
@@ -175,7 +176,7 @@ func (s *Service) GetCustomersByUsageAttribution(ctx context.Context, input cust
 // resolveCustomersByUsageAttribution runs the shared usage-attribution lookup: fetch the raw
 // candidate customers from the adapter, then apply key-over-subject precedence. Both the single-key
 // and bulk service methods route through this so their core resolution stays identical.
-func (s *Service) resolveCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]customer.Customer, error) {
+func (s *Service) resolveCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) {
 	customers, err := s.adapter.GetCustomersByUsageAttribution(ctx, input)
 	if err != nil {
 		return nil, err
@@ -188,13 +189,15 @@ func (s *Service) resolveCustomersByUsageAttribution(ctx context.Context, input
 // either by the customer's own key or by one of its subject keys; when a key matches both a
 // distinct key-owner and a distinct subject-owner, the key-owner takes precedence. Key-over-subject
 // collisions are structurally rare and resolved deterministically here, so they are not surfaced as
-// errors or logs; investigate the underlying data via the database if ever needed. Keys with no
-// match are absent from the returned map.
-func resolveCustomersByKeyWithPrecedence(customers []customer.Customer, keys []string) map[string]customer.Customer {
-	byKey := make(map[string]customer.Customer, len(customers))
-	bySubject := make(map[string]customer.Customer, len(customers))
+// errors or logs; investigate the underlying data via the database if ever needed. Every input key
+// is present in the returned map; keys that match no customer have a nil value.
+func resolveCustomersByKeyWithPrecedence(customers []customer.Customer, keys []string) map[string]*customer.Customer {
+	byKey := make(map[string]*customer.Customer, len(customers))
+	bySubject := make(map[string]*customer.Customer, len(customers))
+
+	for i := range customers {
+		c := &customers[i]
 
-	for _, c := range customers {
 		if c.Key != nil {
 			byKey[*c.Key] = c
 		}
@@ -208,7 +211,7 @@ func resolveCustomersByKeyWithPrecedence(customers []customer.Customer, keys []s
 		}
 	}
 
-	resolved := make(map[string]customer.Customer, len(keys))
+	resolved := make(map[string]*customer.Customer, len(keys))
 
 	for _, k := range keys {
 		if keyOwner, ok := byKey[k]; ok {
@@ -218,7 +221,10 @@ func resolveCustomersByKeyWithPrecedence(customers []customer.Customer, keys []s
 
 		if subjectOwner, ok := bySubject[k]; ok {
 			resolved[k] = subjectOwner
+			continue
 		}
+
+		resolved[k] = nil // present, but no matching customer
 	}
 
 	return resolved

openmeter/customer/service/customer_test.go

@@ -22,19 +22,19 @@ func Test_resolveCustomersByKeyWithPrecedence(t *testing.T) {
 	t.Run("MatchByKey", func(t *testing.T) {
 		resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"key-a"})
 
-		assert.Equal(t, map[string]customer.Customer{"key-a": customerA}, resolved)
+		assert.Equal(t, map[string]*customer.Customer{"key-a": &customerA}, resolved)
 	})
 
 	t.Run("MatchBySubject", func(t *testing.T) {
 		resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"subject-a"})
 
-		assert.Equal(t, map[string]customer.Customer{"subject-a": customerA}, resolved)
+		assert.Equal(t, map[string]*customer.Customer{"subject-a": &customerA}, resolved)
 	})
 
-	t.Run("UnmatchedKeyIsAbsent", func(t *testing.T) {
+	t.Run("UnmatchedKeyIsNil", func(t *testing.T) {
 		resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{customerA}, []string{"no-such-key"})
 
-		assert.Empty(t, resolved)
+		assert.Equal(t, map[string]*customer.Customer{"no-such-key": nil}, resolved)
 	})
 
 	t.Run("KeyOwnerTakesPrecedenceOverDistinctSubjectOwner", func(t *testing.T) {
@@ -55,7 +55,7 @@ func Test_resolveCustomersByKeyWithPrecedence(t *testing.T) {
 			[]string{"shared"},
 		)
 
-		assert.Equal(t, map[string]customer.Customer{"shared": sharedKeyCustomer}, resolved)
+		assert.Equal(t, map[string]*customer.Customer{"shared": &sharedKeyCustomer}, resolved)
 	})
 
 	t.Run("SameCustomerMatchedByOwnKeyAndSubjectKey", func(t *testing.T) {
@@ -69,7 +69,7 @@ func Test_resolveCustomersByKeyWithPrecedence(t *testing.T) {
 
 		resolved := resolveCustomersByKeyWithPrecedence([]customer.Customer{selfMatched}, []string{"dual"})
 
-		assert.Equal(t, map[string]customer.Customer{"dual": selfMatched}, resolved)
+		assert.Equal(t, map[string]*customer.Customer{"dual": &selfMatched}, resolved)
 	})
 
 	t.Run("CrossCollisionResolvesEachKeyToItsOwnKeyOwner", func(t *testing.T) {
@@ -96,9 +96,9 @@ func Test_resolveCustomersByKeyWithPrecedence(t *testing.T) {
 			[]string{"key-1", "key-2"},
 		)
 
-		assert.Equal(t, map[string]customer.Customer{
-			"key-1": crossA,
-			"key-2": crossB,
+		assert.Equal(t, map[string]*customer.Customer{
+			"key-1": &crossA,
+			"key-2": &crossB,
 		}, resolved)
 	})
 }

openmeter/customer/service/customer_usage_attribution_benchmark_test.go

@@ -125,6 +125,8 @@ func BenchmarkCustomersUsageAttributionBulkLookup(b *testing.B) {
 
 	keys := uabench.BulkKeys(customerCount)
 
+	// Count non-nil entries: every input key is present in the result map, with a nil value for
+	// keys that resolved to no customer.
 	get := func(ctx context.Context) (int, error) {
 		resolved, err := svc.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{
 			Namespace: namespace,
@@ -133,7 +135,13 @@ func BenchmarkCustomersUsageAttributionBulkLookup(b *testing.B) {
 		if err != nil {
 			return 0, err
 		}
-		return len(resolved), nil
+		matched := 0
+		for _, c := range resolved {
+			if c != nil {
+				matched++
+			}
+		}
+		return matched, nil
 	}
 
 	count, err := get(b.Context())

openmeter/customer/service/service_test.go

@@ -220,7 +220,7 @@ func Test_CustomerService(t *testing.T) {
 
 // Test_GetCustomersByUsageAttribution protects the bulk method's contract with the same rigor as
 // Test_GetCustomerByUsageAttribution does the single-key one: each input key maps to its customer
-// with key-over-subject precedence, and unmatched keys are simply absent (no error). The predicate
+// with key-over-subject precedence, and unmatched keys are present with a nil value (no error). The predicate
 // itself is covered by the adapter's TestGetCustomersByUsageAttribution, and the fully symmetric
 // {K1,K2} cross-collision by Test_resolveCustomersByKeyWithPrecedence — that state cannot be built
 // through the service here, because CreateCustomer rejects a customer whose key overlaps another
@@ -283,23 +283,25 @@ func Test_GetCustomersByUsageAttribution(t *testing.T) {
 		assert.Equal(t, p.ID, got["shared"].ID, "a direct customer-key match must win over a subject-key match")
 	})
 
-	t.Run("UnmatchedKeysAbsentNoError", func(t *testing.T) {
+	t.Run("UnmatchedKeysNil", func(t *testing.T) {
 		// given:
 		// - a mix of a known key and an unknown one
 		// then:
-		// - the unknown key is simply absent from the map (unlike the single-key method, the bulk
-		//   method does NOT return a not-found error)
+		// - every input key is present in the map; the unknown key has a nil value (unlike the
+		//   single-key method, the bulk method does NOT return a not-found error)
 		known := create(t, "known-key")
 
 		got, err := env.CustomerService.GetCustomersByUsageAttribution(ctx, customer.GetCustomersByUsageAttributionInput{
 			Namespace: namespace,
 			Keys:      []string{"known-key", "totally-unknown"},
 		})
 		require.NoError(t, err)
-		require.Len(t, got, 1)
+		require.Len(t, got, 2)
+		require.NotNil(t, got["known-key"])
 		assert.Equal(t, known.ID, got["known-key"].ID)
-		_, ok := got["totally-unknown"]
-		assert.False(t, ok, "unmatched key must be absent from the result map")
+		unknown, ok := got["totally-unknown"]
+		assert.True(t, ok, "every input key must be present in the result map")
+		assert.Nil(t, unknown, "unmatched key must have a nil value")
 	})
 
 	t.Run("EmptyKeysFailsValidation", func(t *testing.T) {

openmeter/governance/service/service.go

@@ -303,9 +303,9 @@ func (s *service) resolveCustomers(ctx context.Context, input governance.QueryAc
 
 		// Iterate the input keys in order so matched keys and not-found errors keep input ordering.
 		for _, key := range input.CustomerKeys {
-			cus, ok := keyToCustomer[key]
+			cus := keyToCustomer[key]
 
-			if !ok {
+			if cus == nil {
 				queryErrors = append(queryErrors, governance.QueryError{
 					CustomerKey: key,
 					Code:        governance.QueryErrorCustomerNotFound,
@@ -318,7 +318,7 @@ func (s *service) resolveCustomers(ctx context.Context, input governance.QueryAc
 				rc.matched = append(rc.matched, key)
 			} else {
 				customerMap[cus.ID] = &resolvedCustomer{
-					customer: cus,
+					customer: *cus,
 					matched:  []string{key},
 				}
 			}

openmeter/server/server_test.go

@@ -1458,7 +1458,7 @@ func (n NoopCustomerService) GetCustomerByUsageAttribution(ctx context.Context,
 	return &customer.Customer{}, nil
 }
 
-func (n NoopCustomerService) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]customer.Customer, error) {
+func (n NoopCustomerService) GetCustomersByUsageAttribution(ctx context.Context, input customer.GetCustomersByUsageAttributionInput) (map[string]*customer.Customer, error) {
 	return nil, nil
 }