feat: bulk customer usage attribution lookup performance and correctness

e2e/Makefile

@@ -22,6 +22,12 @@ bench-governance-matrix: ## Benchmark governance query latency (full 3x3 custome
 	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_FULL_MATRIX=1 \
 		go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./...
 
+.PHONY: bench-governance-namespace
+bench-governance-namespace: ## Benchmark governance query latency vs. namespace size (fixed 100 customers x 1 feature, 0/10k/100k decoys); requires direct Postgres access
+	$(call print-target)
+	TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) \
+		go test -run='^$$' -bench='^BenchmarkGovernanceQueryNamespaceScale$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./...
+
 .PHONY: test-local
 test-local: ## Run tests against local openmeter
 	$(call print-target)

e2e/governance_bench_namespace_test.go

@@ -0,0 +1,173 @@
+package e2e
+
+import (
+	"fmt"
+	"net/http"
+	"os"
+	"slices"
+	"sort"
+	"strconv"
+	"testing"
+
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/oklog/ulid/v2"
+	"github.com/stretchr/testify/require"
+
+	v3sdk "github.com/openmeterio/openmeter/api/v3/client"
+)
+
+const (
+	defaultGovernanceBenchNamespaceDecoys = 100_000
+	governanceBenchNamespaceDecoysEnv     = "GOV_BENCH_NAMESPACE_DECOYS"
+)
+
+// BenchmarkGovernanceQueryNamespaceScale isolates the effect of total namespace
+// size on governance query latency, independent of the customers/features axes
+// BenchmarkGovernanceQuery already covers. It fixes a small, constant query
+// load (100 customers, matching the OAS customer-key cap, x 1 feature) so the
+// entitlement GetAccess fan-out cost stays flat across sub-benchmarks, and
+// varies only how many other (never-queried) customers/subjects exist in the
+// namespace.
+//
+// This targets the customer usage-attribution resolution path: pre-UNION-ALL,
+// a large decoy count made GetCustomersByUsageAttribution seq-scan the
+// customers table (see #4684 and the follow-up bulk fix); this benchmark shows
+// whether decoy count still moves total request latency post-fix. Decoy
+// counts accumulate across sub-benchmarks (0 -> 10k -> N), so each step only
+// seeds the incremental delta.
+//
+// Decoys are seeded directly via SQL (not HTTP, which would dominate setup
+// time) and require direct Postgres access alongside OPENMETER_ADDRESS (see
+// initE2EPostgresPool). Set GOV_BENCH_NAMESPACE_DECOYS to change the top
+// decoy count from the default 100,000.
+func BenchmarkGovernanceQueryNamespaceScale(b *testing.B) {
+	client := initClient(b)
+	v3 := newV3Client(b)
+	pool := initE2EPostgresPool(b)
+
+	const (
+		queryCustomers = 100
+		queryFeatures  = 1
+	)
+
+	custKeys, featKeys := seedGovernanceFixture(b, client, queryCustomers, queryFeatures)
+	namespace := getCustomerNamespaceByKey(b, pool, custKeys[0])
+	decoyRun := uniqueKey("gov_bench_ns_decoy")
+
+	// Sorted and deduplicated so the cumulative seeding below stays monotonically
+	// increasing regardless of GOV_BENCH_NAMESPACE_DECOYS: an override below the
+	// hardcoded 10_000 tier would otherwise seed 10_000 first, then skip seeding
+	// for the smaller tier (seeded > target) while still reporting that smaller,
+	// now-incorrect decoy count as the benchmark's label.
+	decoyCounts := []int{0, 10_000, governanceBenchNamespaceDecoyCount(b)}
+	sort.Ints(decoyCounts)
+	decoyCounts = slices.Compact(decoyCounts)
+
+	reqBody := v3sdk.GovernanceQueryRequest{
+		Customer: v3sdk.GovernanceQueryRequestCustomers{Keys: custKeys},
+		Feature:  &v3sdk.GovernanceQueryRequestFeatures{Keys: featKeys},
+	}
+
+	seeded := 0
+	for _, target := range decoyCounts {
+		b.Run(fmt.Sprintf("decoys=%d", target), func(b *testing.B) {
+			if target > seeded {
+				seedNamespaceDecoys(b, pool, namespace, decoyRun, seeded, target)
+				seeded = target
+			}
+
+			// Warm-up + correctness gate: a wrong result (e.g. missing customers)
+			// would make the latency number meaningless.
+			resp, err := v3.Governance.QueryAccess(b.Context(), reqBody, v3sdk.GovernanceQueryResultListParams{})
+			v3.requireStatus(http.StatusOK, err)
+			require.Lenf(b, resp.Data, queryCustomers, "expected %d resolved customers", queryCustomers)
+
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				if _, err := v3.Governance.QueryAccess(b.Context(), reqBody, v3sdk.GovernanceQueryResultListParams{}); err != nil {
+					b.Fatalf("governance query failed: %v", err)
+				}
+				if s := v3.statuses.last(); s != http.StatusOK {
+					b.Fatalf("governance query returned %d", s)
+				}
+			}
+			b.StopTimer()
+			b.ReportMetric(float64(target), "decoys")
+		})
+	}
+}
+
+func governanceBenchNamespaceDecoyCount(b *testing.B) int {
+	b.Helper()
+
+	value := os.Getenv(governanceBenchNamespaceDecoysEnv)
+	if value == "" {
+		return defaultGovernanceBenchNamespaceDecoys
+	}
+
+	count, err := strconv.Atoi(value)
+	if err != nil || count < 1 {
+		b.Fatalf("%s must be a positive integer, got %q", governanceBenchNamespaceDecoysEnv, value)
+	}
+
+	return count
+}
+
+func getCustomerNamespaceByKey(tb testing.TB, pool *pgxpool.Pool, key string) string {
+	tb.Helper()
+
+	var namespace string
+	err := pool.QueryRow(
+		tb.Context(),
+		`SELECT namespace FROM customers WHERE key = $1`,
+		key,
+	).Scan(&namespace)
+	require.NoError(tb, err)
+
+	return namespace
+}
+
+// seedNamespaceDecoys bulk-inserts decoy customers (and one subject key each)
+// directly via SQL, numbered (from, to] under the given per-benchmark-run
+// prefix, so repeated calls with a growing `to` only seed the incremental
+// delta, and separate benchmark runs never collide on the same key even if
+// the database was not reset in between. Decoys are never included in a
+// governance query; they exist purely to grow the namespace the resolution
+// query has to search. ANALYZE keeps planner stats current, matching the
+// customer adapter's own usage-attribution benchmark.
+func seedNamespaceDecoys(tb testing.TB, pool *pgxpool.Pool, namespace, runPrefix string, from, to int) {
+	tb.Helper()
+
+	n := to - from
+	ids := make([]string, n)
+	keys := make([]string, n)
+	subjectKeys := make([]string, n)
+
+	for i := range n {
+		ids[i] = ulid.Make().String()
+		keys[i] = fmt.Sprintf("%s_%d", runPrefix, from+i)
+		subjectKeys[i] = keys[i] + "_subj"
+	}
+
+	rctx := tb.Context()
+
+	_, err := pool.Exec(rctx, `
+		INSERT INTO customers (id, namespace, created_at, updated_at, name, key)
+		SELECT id, $1, now(), now(), key, key
+		FROM unnest($2::text[], $3::text[]) AS t(id, key)
+	`, namespace, ids, keys)
+	require.NoError(tb, err)
+
+	_, err = pool.Exec(rctx, `
+		INSERT INTO customer_subjects (namespace, subject_key, created_at, customer_id)
+		SELECT $1, subject_key, now(), customer_id
+		FROM unnest($2::text[], $3::text[]) AS t(subject_key, customer_id)
+	`, namespace, subjectKeys, ids)
+	require.NoError(tb, err)
+
+	_, err = pool.Exec(rctx, "ANALYZE customers")
+	require.NoError(tb, err)
+	_, err = pool.Exec(rctx, "ANALYZE customer_subjects")
+	require.NoError(tb, err)
+}

e2e/helpers.go

@@ -4,14 +4,44 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"os"
+	"strings"
 	"testing"
 
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/samber/lo"
 	"github.com/stretchr/testify/require"
 
 	api "github.com/openmeterio/openmeter/api/client/go"
 )
 
+// initE2EPostgresPool opens a direct Postgres connection alongside OPENMETER_ADDRESS, for
+// e2e checks and fixtures that need to read or seed data the public API doesn't expose
+// (ledger account mappings, bulk namespace seeding). Requires OPENMETER_E2E_POSTGRES_URL,
+// or the local compose stack's default exposed port.
+func initE2EPostgresPool(t testing.TB) *pgxpool.Pool {
+	t.Helper()
+
+	dsn := os.Getenv("OPENMETER_E2E_POSTGRES_URL")
+	if dsn == "" {
+		address := os.Getenv("OPENMETER_ADDRESS")
+		if !strings.Contains(address, "localhost:38888") && !strings.Contains(address, "127.0.0.1:38888") {
+			t.Skipf("this e2e check requires OPENMETER_E2E_POSTGRES_URL or local compose stack at localhost:38888, got %q", address)
+		}
+
+		dsn = "postgres://postgres:postgres@127.0.0.1:35432/postgres?sslmode=disable"
+	}
+
+	pool, err := pgxpool.New(t.Context(), dsn)
+	require.NoError(t, err)
+
+	t.Cleanup(pool.Close)
+
+	require.NoError(t, pool.Ping(t.Context()))
+
+	return pool
+}
+
 // This will not be needed once we get rid of subjects. Cloud middleware already handles subject / customer creation.
 func CreateCustomerWithSubject(t *testing.T, client *api.ClientWithResponses, customerKey string, subjectKey string) *api.Customer {
 	t.Helper()

e2e/ledger_accounts_test.go

@@ -1,7 +1,6 @@
 package e2e
 
 import (
-	"os"
 	"strings"
 	"testing"
 
@@ -38,29 +37,6 @@ func TestNewCustomerHasLedgerAccounts(t *testing.T) {
 	}
 }
 
-func initE2EPostgresPool(t *testing.T) *pgxpool.Pool {
-	t.Helper()
-
-	dsn := os.Getenv("OPENMETER_E2E_POSTGRES_URL")
-	if dsn == "" {
-		address := os.Getenv("OPENMETER_ADDRESS")
-		if !strings.Contains(address, "localhost:38888") && !strings.Contains(address, "127.0.0.1:38888") {
-			t.Skipf("ledger account e2e requires OPENMETER_E2E_POSTGRES_URL or local compose stack at localhost:38888, got %q", address)
-		}
-
-		dsn = "postgres://postgres:postgres@127.0.0.1:35432/postgres?sslmode=disable"
-	}
-
-	pool, err := pgxpool.New(t.Context(), dsn)
-	require.NoError(t, err)
-
-	t.Cleanup(pool.Close)
-
-	require.NoError(t, pool.Ping(t.Context()))
-
-	return pool
-}
-
 func getCustomerNamespace(t *testing.T, pool *pgxpool.Pool, customerID string) string {
 	t.Helper()
 

openmeter/customer/adapter.go

@@ -20,7 +20,6 @@ type CustomerAdapter interface {
 	CreateCustomer(ctx context.Context, params CreateCustomerInput) (*Customer, error)
 	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) ([]Customer, error)
 	UpdateCustomer(ctx context.Context, params UpdateCustomerInput) (*Customer, error)
 }