diff --git a/api/v3/handlers/governance/handler.go b/api/v3/handlers/governance/handler.go new file mode 100644 index 0000000000..afc5b6d0f4 --- /dev/null +++ b/api/v3/handlers/governance/handler.go @@ -0,0 +1,30 @@ +package governance + +import ( + "context" + + "github.com/openmeterio/openmeter/openmeter/governance" + "github.com/openmeterio/openmeter/pkg/framework/transport/httptransport" +) + +type Handler interface { + QueryGovernanceAccess() QueryGovernanceAccessHandler +} + +type handler struct { + resolveNamespace func(ctx context.Context) (string, error) + governanceService governance.Service + options []httptransport.HandlerOption +} + +func New( + resolveNamespace func(ctx context.Context) (string, error), + governanceService governance.Service, + options ...httptransport.HandlerOption, +) Handler { + return &handler{ + resolveNamespace: resolveNamespace, + governanceService: governanceService, + options: options, + } +} diff --git a/api/v3/handlers/governance/mapping.go b/api/v3/handlers/governance/mapping.go new file mode 100644 index 0000000000..0c9e871ea5 --- /dev/null +++ b/api/v3/handlers/governance/mapping.go @@ -0,0 +1,109 @@ +package governance + +import ( + "github.com/oapi-codegen/nullable" + "github.com/samber/lo" + + apiv3 "github.com/openmeterio/openmeter/api/v3" + customershandler "github.com/openmeterio/openmeter/api/v3/handlers/customers" + "github.com/openmeterio/openmeter/openmeter/governance" +) + +// ToAPIGovernanceQueryResponse maps a domain QueryResult to the API response. +func ToAPIGovernanceQueryResponse(res governance.QueryResult, pageSize int) apiv3.GovernanceQueryResponse { + data := make([]apiv3.GovernanceQueryResult, 0, len(res.Customers)) + + for _, c := range res.Customers { + features := make(map[string]apiv3.GovernanceFeatureAccess, len(c.Features)) + for key, fa := range c.Features { + features[key] = toAPIFeatureAccess(fa) + } + + data = append(data, apiv3.GovernanceQueryResult{ + Matched: c.Matched, + Customer: customershandler.ToAPIBillingCustomer(c.Customer), + Features: features, + UpdatedAt: c.UpdatedAt, + }) + } + + errs := make([]apiv3.GovernanceQueryError, 0, len(res.Errors)) + + for _, e := range res.Errors { + errs = append(errs, apiv3.GovernanceQueryError{ + Customer: lo.ToPtr(e.CustomerKey), + Code: toAPIQueryErrorCode(e.Code), + Message: e.Message, + }) + } + + return apiv3.GovernanceQueryResponse{ + Data: data, + Errors: errs, + Meta: toAPICursorMeta(res, pageSize), + } +} + +func toAPIFeatureAccess(fa governance.FeatureAccess) apiv3.GovernanceFeatureAccess { + out := apiv3.GovernanceFeatureAccess{HasAccess: fa.HasAccess} + + if fa.Reason != nil { + out.Reason = &apiv3.GovernanceFeatureAccessReason{ + Code: toAPIReasonCode(fa.Reason.Code), + Message: fa.Reason.Message, + } + } + + return out +} + +func toAPIReasonCode(code governance.ReasonCode) apiv3.GovernanceFeatureAccessReasonCode { + switch code { + case governance.ReasonCodeUsageLimitReached: + return apiv3.GovernanceFeatureAccessReasonCodeUsageLimitReached + case governance.ReasonCodeFeatureUnavailable: + return apiv3.GovernanceFeatureAccessReasonCodeFeatureUnavailable + case governance.ReasonCodeFeatureNotFound: + return apiv3.GovernanceFeatureAccessReasonCodeFeatureNotFound + case governance.ReasonCodeNoCreditAvailable: + return apiv3.GovernanceFeatureAccessReasonCodeNoCreditAvailable + default: + return apiv3.GovernanceFeatureAccessReasonCodeUnknown + } +} + +func toAPIQueryErrorCode(code governance.QueryErrorCode) apiv3.GovernanceQueryErrorCode { + switch code { + case governance.QueryErrorCustomerNotFound: + return apiv3.GovernanceQueryErrorCodeCustomerNotFound + default: + return apiv3.GovernanceQueryErrorCodeUnknown + } +} + +// toAPICursorMeta builds cursor pagination metadata from the domain result. +func toAPICursorMeta(res governance.QueryResult, pageSize int) apiv3.CursorMeta { + meta := apiv3.CursorMeta{ + Page: apiv3.CursorMetaPage{ + Next: nullable.NewNullNullable[string](), + Previous: nullable.NewNullNullable[string](), + Size: float32(pageSize), + }, + } + + if res.First != nil { + meta.Page.First = lo.ToPtr(res.First.Encode()) + if res.HasPrev { + meta.Page.Previous = nullable.NewNullableWithValue(res.First.Encode()) + } + } + + if res.Last != nil { + meta.Page.Last = lo.ToPtr(res.Last.Encode()) + if res.HasNext { + meta.Page.Next = nullable.NewNullableWithValue(res.Last.Encode()) + } + } + + return meta +} diff --git a/api/v3/handlers/governance/query.go b/api/v3/handlers/governance/query.go new file mode 100644 index 0000000000..8662666355 --- /dev/null +++ b/api/v3/handlers/governance/query.go @@ -0,0 +1,140 @@ +package governance + +import ( + "context" + "fmt" + "net/http" + + apiv3 "github.com/openmeterio/openmeter/api/v3" + "github.com/openmeterio/openmeter/api/v3/apierrors" + "github.com/openmeterio/openmeter/openmeter/governance" + "github.com/openmeterio/openmeter/pkg/framework/commonhttp" + "github.com/openmeterio/openmeter/pkg/framework/transport/httptransport" + pagination "github.com/openmeterio/openmeter/pkg/pagination/v2" +) + +const ( + defaultPageSize = 100 + maxPageSize = 100 +) + +type ( + QueryGovernanceAccessParams = apiv3.QueryGovernanceAccessParams + QueryGovernanceAccessResponse = apiv3.GovernanceQueryResponse + QueryGovernanceAccessHandler = httptransport.HandlerWithArgs[governance.QueryAccessInput, QueryGovernanceAccessResponse, QueryGovernanceAccessParams] +) + +func (h *handler) QueryGovernanceAccess() QueryGovernanceAccessHandler { + return httptransport.NewHandlerWithArgs( + func(ctx context.Context, r *http.Request, params QueryGovernanceAccessParams) (governance.QueryAccessInput, error) { + ns, err := h.resolveNamespace(ctx) + if err != nil { + return governance.QueryAccessInput{}, err + } + + var body apiv3.GovernanceQueryRequest + + if err := commonhttp.JSONRequestBodyDecoder(r, &body); err != nil { + return governance.QueryAccessInput{}, err + } + + input := governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: body.Customer.Keys, + PageSize: defaultPageSize, + } + + if body.Feature != nil { + input.FeatureKeys = body.Feature.Keys + } + + if body.IncludeCredits != nil { + input.IncludeCredits = *body.IncludeCredits + } + + if err := applyPaging(ctx, &input, params); err != nil { + return governance.QueryAccessInput{}, err + } + + return input, nil + }, + func(ctx context.Context, input governance.QueryAccessInput) (QueryGovernanceAccessResponse, error) { + res, err := h.governanceService.QueryAccess(ctx, input) + if err != nil { + return QueryGovernanceAccessResponse{}, err + } + + return ToAPIGovernanceQueryResponse(res, input.PageSize), nil + }, + commonhttp.JSONResponseEncoderWithStatus[QueryGovernanceAccessResponse](http.StatusOK), + httptransport.AppendOptions( + h.options, + httptransport.WithOperationName("query-governance-access"), + httptransport.WithErrorEncoder(apierrors.GenericErrorEncoder()), + )..., + ) +} + +// applyPaging parses page[size]/page[after]/page[before] into the service input. +func applyPaging(ctx context.Context, input *governance.QueryAccessInput, params QueryGovernanceAccessParams) error { + if params.Page == nil { + return nil + } + + if params.Page.Size != nil { + if *params.Page.Size < 1 || *params.Page.Size > maxPageSize { + return apierrors.NewBadRequestError(ctx, + fmt.Errorf("page[size] must be between 1 and %d", maxPageSize), + apierrors.InvalidParameters{{ + Field: "page[size]", + Reason: fmt.Sprintf("must be between 1 and %d", maxPageSize), + Source: apierrors.InvalidParamSourceQuery, + }}, + ) + } + + input.PageSize = *params.Page.Size + } + + if params.Page.After != nil && params.Page.Before != nil { + return apierrors.NewBadRequestError(ctx, + fmt.Errorf("page[after] and page[before] are mutually exclusive"), + apierrors.InvalidParameters{{ + Field: "page[after]", + Reason: "cannot be combined with page[before]", + Source: apierrors.InvalidParamSourceQuery, + }}, + ) + } + + if params.Page.After != nil { + cursor, err := decodeCursorParam(ctx, "page[after]", *params.Page.After) + if err != nil { + return err + } + input.After = cursor + } + + if params.Page.Before != nil { + cursor, err := decodeCursorParam(ctx, "page[before]", *params.Page.Before) + if err != nil { + return err + } + input.Before = cursor + } + + return nil +} + +func decodeCursorParam(ctx context.Context, field, raw string) (*pagination.Cursor, error) { + cursor, err := pagination.DecodeCursor(raw) + if err != nil { + return nil, apierrors.NewBadRequestError(ctx, err, apierrors.InvalidParameters{{ + Field: field, + Reason: err.Error(), + Source: apierrors.InvalidParamSourceQuery, + }}) + } + + return cursor, nil +} diff --git a/api/v3/server/routes.go b/api/v3/server/routes.go index e70399f98a..5c36ea179d 100644 --- a/api/v3/server/routes.go +++ b/api/v3/server/routes.go @@ -470,5 +470,5 @@ func (s *Server) UpdateOrganizationDefaultTaxCodes(w http.ResponseWriter, r *htt // Governance func (s *Server) QueryGovernanceAccess(w http.ResponseWriter, r *http.Request, params api.QueryGovernanceAccessParams) { - unimplemented.QueryGovernanceAccess(w, r, params) + s.governanceHandler.QueryGovernanceAccess().With(params).ServeHTTP(w, r) } diff --git a/api/v3/server/server.go b/api/v3/server/server.go index 00d7cbd459..7ab05edb85 100644 --- a/api/v3/server/server.go +++ b/api/v3/server/server.go @@ -27,6 +27,7 @@ import ( eventshandler "github.com/openmeterio/openmeter/api/v3/handlers/events" featurecosthandler "github.com/openmeterio/openmeter/api/v3/handlers/featurecost" featureshandler "github.com/openmeterio/openmeter/api/v3/handlers/features" + governancehandler "github.com/openmeterio/openmeter/api/v3/handlers/governance" llmcosthandler "github.com/openmeterio/openmeter/api/v3/handlers/llmcost" metershandler "github.com/openmeterio/openmeter/api/v3/handlers/meters" planshandler "github.com/openmeterio/openmeter/api/v3/handlers/plans" @@ -46,6 +47,7 @@ import ( "github.com/openmeterio/openmeter/openmeter/currencies" "github.com/openmeterio/openmeter/openmeter/customer" "github.com/openmeterio/openmeter/openmeter/entitlement" + "github.com/openmeterio/openmeter/openmeter/governance" "github.com/openmeterio/openmeter/openmeter/ingest" "github.com/openmeterio/openmeter/openmeter/ledger" "github.com/openmeterio/openmeter/openmeter/ledger/customerbalance" @@ -93,6 +95,7 @@ type Config struct { AccountResolver ledger.AccountResolver CustomerBalanceFacade *customerbalance.Facade EntitlementService entitlement.Service + GovernanceService governance.Service PlanService plan.Service PlanAddonService planaddon.Service PlanSubscriptionService plansubscription.PlanSubscriptionService @@ -159,6 +162,10 @@ func (c *Config) Validate() error { errs = append(errs, errors.New("entitlement service is required")) } + if c.GovernanceService == nil { + errs = append(errs, errors.New("governance service is required")) + } + if c.PlanService == nil { errs = append(errs, errors.New("plan service is required")) } @@ -238,6 +245,7 @@ type Server struct { customersBillingHandler customersbillinghandler.Handler customersCreditsHandler customerscreditshandler.Handler customersEntitlementHandler customersentitlementhandler.Handler + governanceHandler governancehandler.Handler metersHandler metershandler.Handler subscriptionsHandler subscriptionshandler.Handler subscriptionAddonsHandler subscriptionaddonshandler.Handler @@ -318,6 +326,7 @@ func NewServer(config *Config) (*Server, error) { } featuresH := featureshandler.New(resolveNamespace, config.FeatureConnector, config.MeterService, config.LLMCostService, httptransport.WithErrorHandler(config.ErrorHandler)) + governanceHandler := governancehandler.New(resolveNamespace, config.GovernanceService, httptransport.WithErrorHandler(config.ErrorHandler)) var llmcostH llmcosthandler.Handler if config.LLMCostService != nil { @@ -351,6 +360,7 @@ func NewServer(config *Config) (*Server, error) { currenciesHandler: currenciesHandler, featuresHandler: featuresH, featureCostHandler: featureCostH, + governanceHandler: governanceHandler, }, nil } diff --git a/app/common/governance.go b/app/common/governance.go new file mode 100644 index 0000000000..43f83b5726 --- /dev/null +++ b/app/common/governance.go @@ -0,0 +1,31 @@ +package common + +import ( + "github.com/google/wire" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/trace" + + "github.com/openmeterio/openmeter/openmeter/customer" + "github.com/openmeterio/openmeter/openmeter/governance" + governanceservice "github.com/openmeterio/openmeter/openmeter/governance/service" + "github.com/openmeterio/openmeter/openmeter/registry" +) + +var Governance = wire.NewSet( + NewGovernanceService, +) + +func NewGovernanceService( + customer customer.Service, + entitlementRegistry *registry.Entitlement, + tracer trace.Tracer, + meter metric.Meter, +) (governance.Service, error) { + return governanceservice.New(governanceservice.Config{ + Customer: customer, + Entitlement: entitlementRegistry.Entitlement, + Feature: entitlementRegistry.Feature, + Tracer: tracer, + Meter: meter, + }) +} diff --git a/cmd/server/main.go b/cmd/server/main.go index 531ec9f686..fb7b488cb4 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -173,6 +173,7 @@ func main() { EntitlementBalanceConnector: app.EntitlementRegistry.MeteredEntitlement, EntitlementConnector: app.EntitlementRegistry.Entitlement, FeatureConnector: app.FeatureConnector, + GovernanceService: app.GovernanceService, GrantConnector: app.EntitlementRegistry.Grant, GrantRepo: app.EntitlementRegistry.GrantRepo, IngestService: app.IngestService, diff --git a/cmd/server/wire.go b/cmd/server/wire.go index 0741e78f75..61c8f5cbb3 100644 --- a/cmd/server/wire.go +++ b/cmd/server/wire.go @@ -19,6 +19,7 @@ import ( "github.com/openmeterio/openmeter/openmeter/currencies" "github.com/openmeterio/openmeter/openmeter/customer" "github.com/openmeterio/openmeter/openmeter/ent/db" + "github.com/openmeterio/openmeter/openmeter/governance" "github.com/openmeterio/openmeter/openmeter/ingest" "github.com/openmeterio/openmeter/openmeter/ingest/kafkaingest" "github.com/openmeterio/openmeter/openmeter/ledger" @@ -69,6 +70,7 @@ type Application struct { EntitlementRegistry *registry.Entitlement FeatureConnector feature.FeatureConnector FeatureFlags ffx.Service + GovernanceService governance.Service IngestCollector ingest.Collector IngestService ingest.Service KafkaProducer *kafka.Producer @@ -122,6 +124,7 @@ func initializeApplication(ctx context.Context, conf config.Configuration) (Appl common.Entitlement, common.Framework, common.FFX, + common.Governance, common.Kafka, common.KafkaIngest, common.LLMCost, diff --git a/cmd/server/wire_gen.go b/cmd/server/wire_gen.go index 949f2375a1..07b686722d 100644 --- a/cmd/server/wire_gen.go +++ b/cmd/server/wire_gen.go @@ -16,6 +16,7 @@ import ( "github.com/openmeterio/openmeter/openmeter/currencies" "github.com/openmeterio/openmeter/openmeter/customer" "github.com/openmeterio/openmeter/openmeter/ent/db" + "github.com/openmeterio/openmeter/openmeter/governance" "github.com/openmeterio/openmeter/openmeter/ingest" "github.com/openmeterio/openmeter/openmeter/ingest/kafkaingest" "github.com/openmeterio/openmeter/openmeter/ledger" @@ -543,6 +544,17 @@ func initializeApplication(ctx context.Context, conf config.Configuration) (Appl cleanup() return Application{}, nil, err } + governanceService, err := common.NewGovernanceService(customerService, entitlement, tracer, meter) + if err != nil { + cleanup7() + cleanup6() + cleanup5() + cleanup4() + cleanup3() + cleanup2() + cleanup() + return Application{}, nil, err + } dedupeConfiguration := conf.Dedupe producer, err := common.NewKafkaProducer(kafkaIngestConfiguration, logger, commonMetadata) if err != nil { @@ -801,6 +813,7 @@ func initializeApplication(ctx context.Context, conf config.Configuration) (Appl EntitlementRegistry: entitlement, FeatureConnector: featureConnector, FeatureFlags: ffxService, + GovernanceService: governanceService, IngestCollector: ingestCollector, IngestService: ingestService, KafkaProducer: producer, @@ -872,6 +885,7 @@ type Application struct { EntitlementRegistry *registry.Entitlement FeatureConnector feature.FeatureConnector FeatureFlags ffx.Service + GovernanceService governance.Service IngestCollector ingest.Collector IngestService ingest.Service KafkaProducer *kafka2.Producer diff --git a/e2e/Makefile b/e2e/Makefile index aca57447c4..03f9ef7c7f 100644 --- a/e2e/Makefile +++ b/e2e/Makefile @@ -1,5 +1,26 @@ # A Self-Documenting Makefile: http://marmelab.com/blog/2016/02/29/auto-documented-makefile.html +# Address of a running OpenMeter to benchmark against. Defaults to the local stack +# (env-local-up). Override on the CLI: make bench-governance OPENMETER_ADDRESS=... +OPENMETER_ADDRESS ?= http://localhost:38888 +# Iterations per sub-benchmark. Each governance query is sequential; "20x" fixes the +# count so runs are comparable. Bump for tighter means, or use COUNT for variance. +BENCHTIME ?= 20x +# Number of times to repeat the whole benchmark (feed >1 into benchstat for variance). +COUNT ?= 1 + +.PHONY: bench-governance +bench-governance: ## Benchmark governance query latency (1x1 baseline + 10/50/100 diagonal) + $(call print-target) + TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) \ + go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./... + +.PHONY: bench-governance-matrix +bench-governance-matrix: ## Benchmark governance query latency (full 3x3 customers x features matrix) + $(call print-target) + TZ=UTC OPENMETER_ADDRESS=$(OPENMETER_ADDRESS) GOV_BENCH_FULL_MATRIX=1 \ + go test -run='^$$' -bench='^BenchmarkGovernanceQuery$$' -benchmem -benchtime=$(BENCHTIME) -count=$(COUNT) ./... + .PHONY: test-local test-local: ## Run tests against local openmeter $(call print-target) diff --git a/e2e/README.md b/e2e/README.md new file mode 100644 index 0000000000..5ba3d9ffd4 --- /dev/null +++ b/e2e/README.md @@ -0,0 +1,66 @@ +# E2E Tests & Benchmarks + +End-to-end tests and benchmarks that run against a **live OpenMeter stack** over HTTP. +Both skip unless `OPENMETER_ADDRESS` is set. + +## Run the local stack + +```sh +make env-local-up # build + start openmeter + infra (docker compose) +make env-local-down # tear down +``` + +Server listens on `http://localhost:38888`. `config.yaml` is bind-mounted, so config +changes need only `make env-local-up` (force-recreate), not a full down/up. Logs land in +`e2e/logs/`. + +## Benchmarks + +`BenchmarkGovernanceQuery` measures `POST /api/v3/openmeter/governance/query` latency +across a customers × features grid. Seeds boolean entitlements (no usage events). + +```sh +make bench-governance # 1x1 baseline + 10/50/100 diagonal +make bench-governance-matrix # full 3x3 customers x features matrix +``` + +Overridable vars (defaults shown): + +| Var | Default | Purpose | +|---------------------|--------------------------|----------------------------------------------| +| `OPENMETER_ADDRESS` | `http://localhost:38888` | target server | +| `BENCHTIME` | `20x` | iterations per sub-benchmark | +| `COUNT` | `1` | repeat count (use >1 for benchstat variance) | + +For variance / before-after comparison: + +```sh +make bench-governance COUNT=10 | tee baseline.txt +benchstat baseline.txt # mean ± %CV +benchstat baseline.txt after.txt # delta + p-value +``` + +> `ns/op` = mean of sequential request latencies (no concurrency, no tail percentiles). +> Boolean entitlements skip the metered/ClickHouse balance path — a relative algorithmic +> baseline, not a production-latency oracle. + +## Traces (optional) + +OpenMeter runs in a container, so OTLP must target the **host**, not container loopback. +Point the OTLP exporters in `e2e/config.yaml` at `host.docker.internal:4317` and run a +collector on the host (e.g. `grafana/otel-lgtm`). Then query per-size latency percentiles +in Grafana (Tempo, TraceQL metrics): + +``` +{ name = "governance.QueryAccess" && span.customer_key_count = 100 && span.feature_key_count = 100 } + | quantile_over_time(duration, 0.5, 0.9, 0.95, 0.99) +``` + +High-fan-out requests emit thousands of SQL spans; bump Tempo's `max_bytes_per_trace` +(default 5MB) or traces get dropped at ingest. + +## Standard e2e tests + +```sh +make test-local # full down → up → go test ./... → down +``` diff --git a/e2e/governance_bench_test.go b/e2e/governance_bench_test.go new file mode 100644 index 0000000000..27985f8611 --- /dev/null +++ b/e2e/governance_bench_test.go @@ -0,0 +1,161 @@ +package e2e + +import ( + "context" + "fmt" + "net/http" + "os" + "testing" + + "github.com/samber/lo" + "github.com/stretchr/testify/require" + + api "github.com/openmeterio/openmeter/api/client/go" + apiv3 "github.com/openmeterio/openmeter/api/v3" +) + +// BenchmarkGovernanceQuery measures end-to-end latency of +// POST /api/v3/openmeter/governance/query against a running stack +// (OPENMETER_ADDRESS). It is a benchmark, so it only runs under `go test -bench`; +// a plain `go test ./e2e/...` skips it entirely. +// +// What it is good for: a realistic baseline and a before/after signal for +// performance work — it exercises the real router, the OAS layer, real Postgres, +// and the real entitlement GetAccess fan-out over HTTP. +// +// What it is NOT: a production-latency oracle. Entitlements seeded here are +// boolean, so GetAccess skips the metered balance path (ClickHouse). That path is +// the larger production cost and is instrumented separately (entitlement package). +// A metered variant — seeding usage events and waiting for ClickHouse ingestion — +// is the follow-up for measuring it. +// +// Sizes scale customers x features. The customer-count axis drives the GetAccess +// fan-out (the dominant cost); the feature-count axis drives the per-customer +// feature-access map. By default the diagonal (10/50/100% of the 100x100 spec +// ceiling) plus a 1x1 fixed-overhead baseline is run. Set GOV_BENCH_FULL_MATRIX=1 +// to run the full 3x3 matrix, which isolates the two axes. +// +// Seeding is heavy (100x100 = ~10k entitlement creates over HTTP) and runs once +// per sub-benchmark, outside the timed loop. +func BenchmarkGovernanceQuery(b *testing.B) { + client := initClient(b) // skips when OPENMETER_ADDRESS is unset + v3 := newV3Client(b) + + type size struct { + name string + customers int + features int + } + + // Diagonal + fixed-overhead baseline (default). + sizes := []size{ + {"customers=1/features=1", 1, 1}, + {"customers=10/features=10", 10, 10}, + {"customers=50/features=50", 50, 50}, + {"customers=100/features=100", 100, 100}, + } + + // Full 3x3 matrix isolates the customer axis from the feature axis. + if os.Getenv("GOV_BENCH_FULL_MATRIX") != "" { + sizes = nil + sizes = append(sizes, size{"customers=1/features=1", 1, 1}) + for _, c := range []int{10, 50, 100} { + for _, f := range []int{10, 50, 100} { + sizes = append(sizes, size{fmt.Sprintf("customers=%d/features=%d", c, f), c, f}) + } + } + } + + for _, s := range sizes { + b.Run(s.name, func(b *testing.B) { + custKeys, featKeys := seedGovernanceFixture(b, client, s.customers, s.features) + + reqBody := apiv3.GovernanceQueryRequest{ + Customer: apiv3.GovernanceQueryRequestCustomers{Keys: custKeys}, + Feature: &apiv3.GovernanceQueryRequestFeatures{Keys: featKeys}, + } + + // Warm-up + correctness gate: a wrong result (e.g. missing customers) + // would make the latency number meaningless. + status, resp, problem := v3.QueryGovernance(reqBody) + require.Equalf(b, http.StatusOK, status, "governance query failed: %+v", problem) + require.Lenf(b, resp.Data, s.customers, "expected %d resolved customers", s.customers) + require.Lenf(b, resp.Data[0].Features, s.features, "expected %d features per customer", s.features) + + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + status, _, _ := v3.QueryGovernance(reqBody) + if status != http.StatusOK { + b.Fatalf("governance query returned %d", status) + } + } + b.StopTimer() + }) + } +} + +// seedGovernanceFixture creates nFeatures boolean features and nCustomers +// customers, granting each customer a boolean entitlement for every feature +// (nCustomers x nFeatures entitlements). Keys carry a per-run unique prefix so +// repeated runs against the same database do not collide. Returns the customer +// keys and feature keys to query. +func seedGovernanceFixture(b *testing.B, client *api.ClientWithResponses, nCustomers, nFeatures int) (custKeys, featKeys []string) { + b.Helper() + ctx := b.Context() + run := uniqueKey("gov_bench") + + featKeys = make([]string, 0, nFeatures) + for i := 0; i < nFeatures; i++ { + fkey := fmt.Sprintf("%s_feat_%d", run, i) + resp, err := client.CreateFeatureWithResponse(ctx, api.CreateFeatureJSONRequestBody{ + Key: fkey, + Name: fkey, + }) + require.NoError(b, err) + require.Equalf(b, http.StatusCreated, resp.StatusCode(), "create feature: %s", resp.Body) + featKeys = append(featKeys, fkey) + } + + custKeys = make([]string, 0, nCustomers) + for c := 0; c < nCustomers; c++ { + ckey := fmt.Sprintf("%s_cust_%d", run, c) + skey := ckey + "_subj" + + subResp, err := client.UpsertSubjectWithResponse(ctx, api.UpsertSubjectJSONRequestBody{api.SubjectUpsert{Key: skey}}) + require.NoError(b, err) + require.Equalf(b, http.StatusOK, subResp.StatusCode(), "upsert subject: %s", subResp.Body) + + custResp, err := client.CreateCustomerWithResponse(ctx, api.CreateCustomerJSONRequestBody{ + Key: lo.ToPtr(ckey), + Name: ckey, + UsageAttribution: &api.CustomerUsageAttribution{SubjectKeys: []string{skey}}, + }) + require.NoError(b, err) + require.Equalf(b, http.StatusCreated, custResp.StatusCode(), "create customer: %s", custResp.Body) + custID := custResp.JSON201.Id + + for _, fkey := range featKeys { + grantBooleanEntitlement(b, client, ctx, custID, fkey) + } + custKeys = append(custKeys, ckey) + } + + return custKeys, featKeys +} + +// grantBooleanEntitlement creates a boolean entitlement for the given customer and +// feature key via the V2 customer-entitlement endpoint. +func grantBooleanEntitlement(b *testing.B, client *api.ClientWithResponses, ctx context.Context, custID, featureKey string) { + b.Helper() + + var body api.CreateCustomerEntitlementV2JSONRequestBody + require.NoError(b, body.FromEntitlementBooleanCreateInputs(api.EntitlementBooleanCreateInputs{ + Type: api.EntitlementBooleanCreateInputsTypeBoolean, + FeatureKey: lo.ToPtr(featureKey), + })) + + resp, err := client.CreateCustomerEntitlementV2WithResponse(ctx, custID, body) + require.NoError(b, err) + require.Equalf(b, http.StatusCreated, resp.StatusCode(), "create boolean entitlement: %s", resp.Body) +} diff --git a/e2e/setup_test.go b/e2e/setup_test.go index f0773d814b..8d4e8af090 100644 --- a/e2e/setup_test.go +++ b/e2e/setup_test.go @@ -9,7 +9,7 @@ import ( api "github.com/openmeterio/openmeter/api/client/go" ) -func initClient(t *testing.T) *api.ClientWithResponses { +func initClient(t testing.TB) *api.ClientWithResponses { t.Helper() address := os.Getenv("OPENMETER_ADDRESS") diff --git a/e2e/v3helpers_test.go b/e2e/v3helpers_test.go index 90fba3e8ce..1621dff514 100644 --- a/e2e/v3helpers_test.go +++ b/e2e/v3helpers_test.go @@ -36,13 +36,13 @@ const v3RequestTimeout = 30 * time.Second // collide, and subtests are written to be independent rather than rely on // ordering. type v3Client struct { - t *testing.T + t testing.TB baseURL string } // newV3Client returns a client pointed at $OPENMETER_ADDRESS. Skips the test // when the variable is unset. -func newV3Client(t *testing.T) *v3Client { +func newV3Client(t testing.TB) *v3Client { t.Helper() address := os.Getenv("OPENMETER_ADDRESS") @@ -176,6 +176,15 @@ func (c *v3Client) CreateFeature(body apiv3.CreateFeatureRequest) (int, *apiv3.F return decodeTyped[apiv3.Feature](c, status, raw, problem, http.StatusCreated) } +// --- Governance --- + +// QueryGovernance posts a governance access query. Page size defaults server-side +// to the max (100), so callers querying <=100 customers receive a single page. +func (c *v3Client) QueryGovernance(body apiv3.GovernanceQueryRequest) (int, *apiv3.GovernanceQueryResponse, *v3Problem) { + status, raw, problem := c.do(http.MethodPost, "/governance/query", body) + return decodeTyped[apiv3.GovernanceQueryResponse](c, status, raw, problem, http.StatusOK) +} + // --- Plans --- func (c *v3Client) CreatePlan(body apiv3.CreatePlanRequest) (int, *apiv3.BillingPlan, *v3Problem) { diff --git a/openmeter/governance/governance.go b/openmeter/governance/governance.go new file mode 100644 index 0000000000..5b163a5101 --- /dev/null +++ b/openmeter/governance/governance.go @@ -0,0 +1,136 @@ +package governance + +import ( + "errors" + "time" + + "github.com/openmeterio/openmeter/openmeter/customer" + "github.com/openmeterio/openmeter/pkg/models" + pagination "github.com/openmeterio/openmeter/pkg/pagination/v2" +) + +// ReasonCode is the machine-readable reason a customer does not have access to a feature. +type ReasonCode string + +const ( + ReasonCodeUnknown ReasonCode = "unknown" + ReasonCodeUsageLimitReached ReasonCode = "usage_limit_reached" + ReasonCodeFeatureUnavailable ReasonCode = "feature_unavailable" + ReasonCodeFeatureNotFound ReasonCode = "feature_not_found" + ReasonCodeNoCreditAvailable ReasonCode = "no_credit_available" +) + +var AccessReasonUsageLimitReached = &AccessReason{ + Code: ReasonCodeUsageLimitReached, + Message: "usage limit for feature reached", +} + +var AccessReasonFeatureUnavailable = &AccessReason{ + Code: ReasonCodeFeatureUnavailable, + Message: "feature is not available for customer", +} + +var AccessReasonFeatureNotFound = &AccessReason{ + Code: ReasonCodeFeatureNotFound, + Message: "feature is not found", +} + +// AccessReason explains why a feature is not accessible. +type AccessReason struct { + Code ReasonCode + Message string +} + +// FeatureAccess is the access status for a single feature. +type FeatureAccess struct { + HasAccess bool + // Reason is set when HasAccess is false. + Reason *AccessReason +} + +// CustomerAccess is the access evaluation for a single resolved customer. +type CustomerAccess struct { + // Customer the matched identifiers resolved to. + Customer customer.Customer + // Matched lists the request identifiers (customer key or usage-attribution subject + // key) that resolved to this customer. + Matched []string + // Features maps feature key to its access status. + Features map[string]FeatureAccess + // UpdatedAt is the time the access state was evaluated. + UpdatedAt time.Time +} + +// QueryErrorCode is the machine-readable code for a per-customer query error. +type QueryErrorCode string + +const ( + QueryErrorUnknown QueryErrorCode = "unknown" + QueryErrorCustomerNotFound QueryErrorCode = "customer_not_found" +) + +// QueryError is a partial error for a single input identifier. +type QueryError struct { + // CustomerKey is the request identifier that produced this error. + CustomerKey string + Code QueryErrorCode + Message string +} + +var _ models.Validator = (*QueryAccessInput)(nil) + +// QueryAccessInput is the input for evaluating governance access. +type QueryAccessInput struct { + Namespace string + // CustomerKeys are arbitrary identifiers — each a customer key or a usage-attribution + // subject key. Identifiers that cannot be resolved are reported in QueryResult.Errors. + CustomerKeys []string + // FeatureKeys, when non-empty, restricts evaluation to those feature keys. When empty, + // every non-archived feature in the namespace is evaluated. + FeatureKeys []string + // IncludeCredits requests credit-balance evaluation. Not yet implemented. + IncludeCredits bool + + // Pagination over the resolved customers (sorted by CreatedAt, ID). At most one of + // After/Before may be set. + PageSize int + After *pagination.Cursor + Before *pagination.Cursor +} + +func (i QueryAccessInput) Validate() error { + var errs []error + + if i.Namespace == "" { + errs = append(errs, errors.New("namespace is required")) + } + + if len(i.CustomerKeys) == 0 { + errs = append(errs, errors.New("at least one customer key is required")) + } + + if i.PageSize < 1 { + errs = append(errs, errors.New("page size must be positive")) + } + + if i.After != nil && i.Before != nil { + errs = append(errs, errors.New("after and before cursors are mutually exclusive")) + } + + return models.NewNillableGenericValidationError(errors.Join(errs...)) +} + +// QueryResult is the paged result of a governance access query. +type QueryResult struct { + // Customers are the access evaluations for the current page, ordered by (CreatedAt, ID). + Customers []CustomerAccess + // Errors are partial errors for unresolved input identifiers. + Errors []QueryError + + // HasPrev/HasNext indicate adjacent pages relative to the current one. + HasPrev bool + HasNext bool + // First/Last are the cursors of the first and last item on the current page. + First *pagination.Cursor + Last *pagination.Cursor +} diff --git a/openmeter/governance/service.go b/openmeter/governance/service.go new file mode 100644 index 0000000000..0bc073531e --- /dev/null +++ b/openmeter/governance/service.go @@ -0,0 +1,9 @@ +package governance + +import "context" + +// Service evaluates feature access for customers by composing the customer, entitlement, +// and feature services. It owns no persistence of its own. +type Service interface { + QueryAccess(ctx context.Context, input QueryAccessInput) (QueryResult, error) +} diff --git a/openmeter/governance/service/mapping.go b/openmeter/governance/service/mapping.go new file mode 100644 index 0000000000..b64a3afff8 --- /dev/null +++ b/openmeter/governance/service/mapping.go @@ -0,0 +1,52 @@ +package service + +import ( + "github.com/openmeterio/openmeter/openmeter/entitlement" + booleanentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/boolean" + meteredentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/metered" + staticentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/static" + "github.com/openmeterio/openmeter/openmeter/governance" +) + +// mapEntitlementToAccess converts an entitlement value to a governance feature access result. +// When HasAccess is false, the reason code is derived from the entitlement type. +func mapEntitlementToAccess(v entitlement.EntitlementValue) governance.FeatureAccess { + switch ent := v.(type) { + case *meteredentitlement.MeteredEntitlementValue: + if ent.HasAccess() { + return governance.FeatureAccess{HasAccess: true} + } + + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonUsageLimitReached, + } + + case *booleanentitlement.BooleanEntitlementValue: + if ent.HasAccess() { + return governance.FeatureAccess{HasAccess: true} + } + + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureUnavailable, + } + + case *staticentitlement.StaticEntitlementValue: + if ent.HasAccess() { + return governance.FeatureAccess{HasAccess: true} + } + + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureUnavailable, + } + + default: + // NoAccessValue or unknown type + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureUnavailable, + } + } +} diff --git a/openmeter/governance/service/mapping_test.go b/openmeter/governance/service/mapping_test.go new file mode 100644 index 0000000000..31d3e2b552 --- /dev/null +++ b/openmeter/governance/service/mapping_test.go @@ -0,0 +1,71 @@ +package service + +import ( + "testing" + + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + + "github.com/openmeterio/openmeter/openmeter/entitlement" + booleanentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/boolean" + meteredentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/metered" + staticentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/static" + "github.com/openmeterio/openmeter/openmeter/governance" +) + +func TestMapEntitlementToAccess(t *testing.T) { + tests := []struct { + name string + value entitlement.EntitlementValue + wantHasAccess bool + wantCode *governance.ReasonCode + }{ + { + name: "metered with balance — has access", + value: &meteredentitlement.MeteredEntitlementValue{Balance: 10}, + wantHasAccess: true, + }, + { + name: "metered exhausted — usage limit reached", + value: &meteredentitlement.MeteredEntitlementValue{Balance: 0}, + wantHasAccess: false, + wantCode: lo.ToPtr(governance.ReasonCodeUsageLimitReached), + }, + { + // BooleanEntitlementValue is always HasAccess=true; the gateway returns + // NoAccessValue when the entitlement is inactive/not in plan. + name: "boolean — has access", + value: &booleanentitlement.BooleanEntitlementValue{}, + wantHasAccess: true, + }, + { + // StaticEntitlementValue is always HasAccess=true. + name: "static — has access", + value: &staticentitlement.StaticEntitlementValue{Config: `{"limit":100}`}, + wantHasAccess: true, + }, + { + // NoAccessValue is returned when the entitlement is inactive (not in current period). + name: "no access value — feature unavailable", + value: &entitlement.NoAccessValue{}, + wantHasAccess: false, + wantCode: lo.ToPtr(governance.ReasonCodeFeatureUnavailable), + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + got := mapEntitlementToAccess(tc.value) + + assert.Equal(t, tc.wantHasAccess, got.HasAccess) + + if tc.wantCode != nil { + if assert.NotNil(t, got.Reason) { + assert.Equal(t, *tc.wantCode, got.Reason.Code) + } + } else { + assert.Nil(t, got.Reason) + } + }) + } +} diff --git a/openmeter/governance/service/service.go b/openmeter/governance/service/service.go new file mode 100644 index 0000000000..33e25ef823 --- /dev/null +++ b/openmeter/governance/service/service.go @@ -0,0 +1,609 @@ +package service + +import ( + "context" + "errors" + "fmt" + "sort" + "time" + + "github.com/samber/lo" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" + "go.opentelemetry.io/otel/trace" + + "github.com/openmeterio/openmeter/openmeter/customer" + "github.com/openmeterio/openmeter/openmeter/entitlement" + "github.com/openmeterio/openmeter/openmeter/governance" + "github.com/openmeterio/openmeter/openmeter/productcatalog/feature" + "github.com/openmeterio/openmeter/pkg/clock" + "github.com/openmeterio/openmeter/pkg/framework/tracex" + "github.com/openmeterio/openmeter/pkg/models" + pagination "github.com/openmeterio/openmeter/pkg/pagination/v2" +) + +// featureFetchLimit caps the org-wide feature fetch used when no feature filter is given. +// Acceptable for prototype scale; revisit if feature counts grow large. +const featureFetchLimit = 10_000 + +// Config holds the collaborating services for the governance Service. +type Config struct { + Customer customer.Service + Entitlement entitlement.Service + Feature feature.FeatureConnector + Tracer trace.Tracer + Meter metric.Meter +} + +func (c Config) Validate() error { + var errs []error + + if c.Customer == nil { + errs = append(errs, errors.New("customer service is required")) + } + + if c.Entitlement == nil { + errs = append(errs, errors.New("entitlement service is required")) + } + + if c.Feature == nil { + errs = append(errs, errors.New("feature connector is required")) + } + + if c.Tracer == nil { + errs = append(errs, errors.New("tracer is required")) + } + + if c.Meter == nil { + errs = append(errs, errors.New("meter is required")) + } + + return errors.Join(errs...) +} + +func New(config Config) (governance.Service, error) { + if err := config.Validate(); err != nil { + return nil, err + } + + metrics, err := newMetrics(config.Meter) + if err != nil { + return nil, err + } + + return &service{ + customerService: config.Customer, + entitlementService: config.Entitlement, + featureConnector: config.Feature, + tracer: config.Tracer, + metrics: metrics, + }, nil +} + +type service struct { + customerService customer.Service + entitlementService entitlement.Service + featureConnector feature.FeatureConnector + tracer trace.Tracer + metrics queryMetrics +} + +// queryMetrics holds the instruments for the governance query endpoint. These are unsampled +// (unlike spans), so they back SLOs, alerting, and capacity dashboards. Per-request counts +// are recorded as histogram observations rather than attributes to keep series cardinality +// bounded; only low-cardinality enums are used as counter attributes. +type queryMetrics struct { + // requests counts queries, broken down by pagination direction and whether the + // all-org-features (no filter) path was taken. + requests metric.Int64Counter + // customersNotFound counts input keys that did not resolve to a customer. These return + // HTTP 200 with a partial error, so they are invisible to HTTP-level metrics. + customersNotFound metric.Int64Counter + // featureAccess records the number of feature evaluations per query — the work unit that + // drives latency (~customers × features). + featureAccess metric.Int64Histogram + // customerKeys records the number of customer keys requested per query. + customerKeys metric.Int64Histogram +} + +func newMetrics(meter metric.Meter) (queryMetrics, error) { + requests, err := meter.Int64Counter( + "openmeter.governance.query.requests", + metric.WithDescription("Number of governance access queries"), + metric.WithUnit("{request}"), + ) + if err != nil { + return queryMetrics{}, fmt.Errorf("failed to create requests counter: %w", err) + } + + customersNotFound, err := meter.Int64Counter( + "openmeter.governance.query.customers_not_found", + metric.WithDescription("Number of customer keys that did not resolve to a customer"), + metric.WithUnit("{customer}"), + ) + if err != nil { + return queryMetrics{}, fmt.Errorf("failed to create customers_not_found counter: %w", err) + } + + featureAccess, err := meter.Int64Histogram( + "openmeter.governance.query.feature_access", + metric.WithDescription("Number of feature evaluations per governance query"), + metric.WithUnit("{evaluation}"), + ) + if err != nil { + return queryMetrics{}, fmt.Errorf("failed to create feature_access histogram: %w", err) + } + + customerKeys, err := meter.Int64Histogram( + "openmeter.governance.query.customer_keys", + metric.WithDescription("Number of customer keys requested per governance query"), + metric.WithUnit("{key}"), + ) + if err != nil { + return queryMetrics{}, fmt.Errorf("failed to create customer_keys histogram: %w", err) + } + + return queryMetrics{ + requests: requests, + customersNotFound: customersNotFound, + featureAccess: featureAccess, + customerKeys: customerKeys, + }, nil +} + +var _ governance.Service = (*service)(nil) + +// QueryAccess evaluates feature access for a caller-supplied set of customer keys. +// +// Pagination is in-memory, by design. The reason is the BOUND: the input is OAS-capped at +// 100 keys (@maxItems), so the resolvable set is ≤100 — sorting and slicing it in memory is +// trivially cheap, and pushing a keyset cursor into the DB would buy nothing. (It is NOT +// because there's no collection to paginate: customer.ListCustomers is a DB-side, orderable, +// paginated query and could resolve+order+limit the key set via a `Key $in [...] $or +// usageAttributionSubjectKey $in [...]` filter. That path only becomes worthwhile if the +// 100-key cap is ever lifted — see below.) +// +// Phase order and what each page actually costs (e.g. page size 10 over 100 resolved): +// 1. resolveCustomers resolves ALL keys (today: a point lookup per key) — runs in full +// regardless of page size, because the sort key is (CreatedAt, ID), customer fields: +// page order can't be established without first resolving every key to a customer. +// 2. sort the full set in memory. +// 3. paginate slices to PageSize (10). +// 4. resolveAccess runs only over the page (10): the expensive per-customer GetAccess +// fan-out + one listOrgFeatures. So the dominant cost IS page-limited; only the cheaper +// full-set resolution is paid in full (and repeated per page across a paging client). +// +// Possible optimization (orthogonal to pagination): replace the per-key point lookups in +// resolveCustomers with a single customer.ListCustomers call using an `$in` key filter +// (N lookups → 1 query). Dedup, the per-input `matched` mapping, and per-key not-found +// reporting stay in memory either way, since the input mixes customer keys and subject keys. +// +// Pagination would only need to move into the adapter (as a keyset query) if the contract +// gained an unbounded mode — "all customers in a namespace", or dropping the 100-key cap. +// Not the case today. +func (s *service) QueryAccess(ctx context.Context, input governance.QueryAccessInput) (governance.QueryResult, error) { + fn := func(ctx context.Context) (governance.QueryResult, error) { + if err := input.Validate(); err != nil { + return governance.QueryResult{}, err + } + + span := trace.SpanFromContext(ctx) + + span.SetAttributes( + attribute.String("namespace", input.Namespace), + attribute.Int("customer_key_count", len(input.CustomerKeys)), + attribute.Int("feature_key_count", len(input.FeatureKeys)), + attribute.Int("page.size", input.PageSize), + attribute.String("direction", paginationDirection(input)), + ) + + customers, err := s.resolveCustomers(ctx, input) + if err != nil { + return governance.QueryResult{}, err + } + + // Sort by (CreatedAt, ID) for stable cursor pagination. + sortedCustomers := lo.Values(customers.resolvedCustomers) + + sort.Slice(sortedCustomers, func(i, j int) bool { + ti := sortedCustomers[i].customer.CreatedAt + tj := sortedCustomers[j].customer.CreatedAt + + if !ti.Equal(tj) { + return ti.Before(tj) + } + + return sortedCustomers[i].customer.ID < sortedCustomers[j].customer.ID + }) + + paginatedCustomers := paginate(sortedCustomers, input) + + results, err := s.resolveAccess(ctx, input, paginatedCustomers.customers) + if err != nil { + return governance.QueryResult{}, err + } + + out := governance.QueryResult{ + Customers: results, + Errors: customers.queryErrors, + HasPrev: paginatedCustomers.hasPrev, + HasNext: paginatedCustomers.hasNext, + } + + if len(paginatedCustomers.customers) > 0 { + out.First = lo.ToPtr(cursorFor(paginatedCustomers.customers[0])) + out.Last = lo.ToPtr(cursorFor(paginatedCustomers.customers[len(paginatedCustomers.customers)-1])) + } + + s.recordQueryMetrics(ctx, input, out, len(customers.queryErrors)) + + return out, nil + } + + return tracex.Start[governance.QueryResult](ctx, s.tracer, "governance.QueryAccess").Wrap(fn) +} + +// recordQueryMetrics emits the unsampled query metrics. namespace is a per-tenant label +// (consistent with ingest/sink/balanceworker); per-request counts are histogram values, and +// only low-cardinality enums (direction, all_org_features) are counter attributes. +func (s *service) recordQueryMetrics(ctx context.Context, input governance.QueryAccessInput, out governance.QueryResult, notFound int) { + namespaceAttr := attribute.String("namespace", input.Namespace) + + s.metrics.requests.Add(ctx, 1, metric.WithAttributes( + namespaceAttr, + attribute.String("direction", paginationDirection(input)), + attribute.Bool("all_org_features", len(input.FeatureKeys) == 0), + )) + + if notFound > 0 { + s.metrics.customersNotFound.Add(ctx, int64(notFound), metric.WithAttributes(namespaceAttr)) + } + + featureAccessTotal := lo.SumBy(out.Customers, func(c governance.CustomerAccess) int { + return len(c.Features) + }) + + s.metrics.featureAccess.Record(ctx, int64(featureAccessTotal), metric.WithAttributes(namespaceAttr)) + s.metrics.customerKeys.Record(ctx, int64(len(input.CustomerKeys)), metric.WithAttributes(namespaceAttr)) +} + +// resolvedCustomer groups the matched input keys for a single customer. +type resolvedCustomer struct { + customer customer.Customer + matched []string +} + +type resolveCustomersResult struct { + resolvedCustomers map[string]*resolvedCustomer + queryErrors []governance.QueryError +} + +// resolveCustomers resolves each input key to a customer, deduplicating by customer ID. +// Keys that resolve to no customer are collected as customer-not-found query errors rather +// than failing the whole request. +func (s *service) resolveCustomers(ctx context.Context, input governance.QueryAccessInput) (resolveCustomersResult, error) { + fn := func(ctx context.Context) (resolveCustomersResult, error) { + span := trace.SpanFromContext(ctx) + + span.SetAttributes( + attribute.Int("requested", len(input.CustomerKeys)), + ) + + customerMap := make(map[string]*resolvedCustomer) + var queryErrors []governance.QueryError + + for _, key := range input.CustomerKeys { + cus, err := s.customerService.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{ + Namespace: input.Namespace, + Key: key, + }) + if err != nil { + if models.IsGenericNotFoundError(err) { + queryErrors = append(queryErrors, governance.QueryError{ + CustomerKey: key, + Code: governance.QueryErrorCustomerNotFound, + Message: "customer not found", + }) + continue + } + return resolveCustomersResult{}, fmt.Errorf("failed to resolve customer key %q: %w", key, err) + } + + if rc, ok := customerMap[cus.ID]; ok { + rc.matched = append(rc.matched, key) + } else { + customerMap[cus.ID] = &resolvedCustomer{ + customer: *cus, + matched: []string{key}, + } + } + } + + span.SetAttributes( + attribute.Int("resolved", len(customerMap)), + attribute.Int("not_found", len(queryErrors)), + ) + + return resolveCustomersResult{ + resolvedCustomers: customerMap, + queryErrors: queryErrors, + }, nil + } + + return tracex.Start[resolveCustomersResult](ctx, s.tracer, "governance.resolveCustomers").Wrap(fn) +} + +// resolveAccess resolves entitlement access for each customer on the current page and maps it +// to feature access. UpdatedAt is stamped once for the whole page. +// +// When no feature filter is given, the org-wide feature list is namespace-scoped, so it is +// fetched once for the whole page rather than per customer. +func (s *service) resolveAccess(ctx context.Context, input governance.QueryAccessInput, customers []*resolvedCustomer) ([]governance.CustomerAccess, error) { + fn := func(ctx context.Context) ([]governance.CustomerAccess, error) { + allOrgFeatures := len(input.FeatureKeys) == 0 + + span := trace.SpanFromContext(ctx) + + span.SetAttributes( + attribute.Int("customer_count", len(customers)), + attribute.Int("feature_filter_count", len(input.FeatureKeys)), + attribute.Bool("all_org_features", allOrgFeatures), + ) + + // On the all-org path, fetch the namespace-wide feature list once for the whole page. + var orgFeatures []feature.Feature + + if allOrgFeatures && len(customers) > 0 { + var err error + + orgFeatures, err = s.listOrgFeatures(ctx, input.Namespace) + if err != nil { + return nil, err + } + + span.SetAttributes( + attribute.Int("org_feature_count", len(orgFeatures)), + ) + } + + now := clock.Now() + results := make([]governance.CustomerAccess, 0, len(customers)) + absentFeatureLookups := 0 + featureAccessTotal := 0 + + for _, rc := range customers { + access, err := s.entitlementService.GetAccess(ctx, input.Namespace, rc.customer.ID) + if err != nil { + return nil, fmt.Errorf("failed to get access for customer %s: %w", rc.customer.ID, err) + } + + featureAccessResult, err := s.buildFeatureAccess(ctx, input.Namespace, input.FeatureKeys, orgFeatures, access) + if err != nil { + return nil, fmt.Errorf("failed to build feature access for customer %s: %w", rc.customer.ID, err) + } + + absentFeatureLookups += featureAccessResult.absentLookups + featureAccessTotal += len(featureAccessResult.featureAccess) + + results = append(results, governance.CustomerAccess{ + Customer: rc.customer, + Matched: rc.matched, + Features: featureAccessResult.featureAccess, + UpdatedAt: now, + }) + } + + span.SetAttributes( + attribute.Int("absent_feature_lookups", absentFeatureLookups), + attribute.Int("feature_access_total", featureAccessTotal), + ) + + return results, nil + } + + return tracex.Start[[]governance.CustomerAccess](ctx, s.tracer, "governance.resolveAccess").Wrap(fn) +} + +// paginationDirection reports the pagination mode for span attribution. +func paginationDirection(input governance.QueryAccessInput) string { + switch { + case input.Before != nil: + return "before" + case input.After != nil: + return "after" + default: + return "first" + } +} + +// cursorFor builds the pagination cursor for a resolved customer. CreatedAt is truncated +// to second precision to match the RFC3339 encoding used by cursor strings. +func cursorFor(rc *resolvedCustomer) pagination.Cursor { + return pagination.NewCursor(rc.customer.CreatedAt.Truncate(time.Second), rc.customer.ID) +} + +type paginationResult struct { + customers []*resolvedCustomer + hasPrev bool + hasNext bool +} + +// paginate applies cursor pagination over the sorted customers and reports whether adjacent +// pages exist. Exactly one of input.After / input.Before may be set (enforced by Validate). +func paginate(customers []*resolvedCustomer, input governance.QueryAccessInput) paginationResult { + if input.Before != nil { + // Backward: take the last pageSize items strictly before the cursor. + bc := *input.Before + end := 0 + + for i, rc := range customers { + c := cursorFor(rc) + + if c.Time.After(bc.Time) || (c.Time.Equal(bc.Time) && c.ID >= bc.ID) { + break + } + + end = i + 1 + } + + candidates := customers[:end] + hasPrev := len(candidates) > input.PageSize + + if hasPrev { + candidates = candidates[len(candidates)-input.PageSize:] + } + + // next is always set in backward mode: the before-cursor item itself is forward. + return paginationResult{ + customers: candidates, + hasPrev: hasPrev, + hasNext: true, + } + } + + // Forward (after cursor or first page). + start := 0 + + if input.After != nil { + ac := *input.After + start = len(customers) // beyond all items if cursor is past the end + + for i, rc := range customers { + c := cursorFor(rc) + + if c.Time.After(ac.Time) || (c.Time.Equal(ac.Time) && c.ID > ac.ID) { + start = i + break + } + } + } + + hasPrev := start > 0 + page := customers[start:] + hasNext := len(page) > input.PageSize + + if hasNext { + page = page[:input.PageSize] + } + + return paginationResult{ + customers: page, + hasPrev: hasPrev, + hasNext: hasNext, + } +} + +type buildFeatureAccessResult struct { + featureAccess map[string]governance.FeatureAccess + absentLookups int +} + +// buildFeatureAccess returns the feature access map for a single customer, along with the +// number of absent-feature lookups it performed (per-feature GetFeature calls), for span +// attribution. +// +// If featureKeys is non-empty, only those keys are evaluated. If featureKeys is empty, the +// pre-fetched orgFeatures slice (namespace-wide, resolved once by the caller) is used; features +// the customer has no entitlement for are marked feature-unavailable. +func (s *service) buildFeatureAccess(ctx context.Context, ns string, featureKeys []string, orgFeatures []feature.Feature, access entitlement.Access) (buildFeatureAccessResult, error) { + result := make(map[string]governance.FeatureAccess) + + if len(featureKeys) == 0 { + for _, f := range orgFeatures { + if ev, ok := access.Entitlements[f.Key]; ok { + result[f.Key] = mapEntitlementToAccess(ev.Value) + } else { + result[f.Key] = governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureUnavailable, + } + } + } + + return buildFeatureAccessResult{ + featureAccess: result, + absentLookups: 0, + }, nil + } + + absentLookups := 0 + + for _, key := range featureKeys { + ev, ok := access.Entitlements[key] + + if !ok { + absentLookups++ + + fa, err := s.resolveAbsentFeature(ctx, ns, key) + if err != nil { + return buildFeatureAccessResult{ + featureAccess: nil, + absentLookups: absentLookups, + }, err + } + + result[key] = fa + + continue + } + + result[key] = mapEntitlementToAccess(ev.Value) + } + + return buildFeatureAccessResult{ + featureAccess: result, + absentLookups: absentLookups, + }, nil +} + +// listOrgFeatures fetches all non-archived features in the namespace in one shot. +func (s *service) listOrgFeatures(ctx context.Context, ns string) ([]feature.Feature, error) { + fn := func(ctx context.Context) ([]feature.Feature, error) { + span := trace.SpanFromContext(ctx) + + span.SetAttributes( + attribute.Int("limit", featureFetchLimit), + ) + + res, err := s.featureConnector.ListFeatures(ctx, feature.ListFeaturesParams{ + Namespace: ns, + IncludeArchived: false, + Limit: featureFetchLimit, + }) + if err != nil { + return nil, fmt.Errorf("failed to list org features: %w", err) + } + + span.SetAttributes( + attribute.Int("feature_count", len(res.Items)), + ) + + return res.Items, nil + } + + return tracex.Start[[]feature.Feature](ctx, s.tracer, "governance.listOrgFeatures").Wrap(fn) +} + +// resolveAbsentFeature determines why a requested feature key is absent from GetAccess results: +// either the feature doesn't exist in the org (feature-not-found) or the customer has no +// entitlement for it (feature-unavailable). +func (s *service) resolveAbsentFeature(ctx context.Context, ns, featureKey string) (governance.FeatureAccess, error) { + _, err := s.featureConnector.GetFeature(ctx, ns, featureKey, feature.IncludeArchivedFeatureFalse) + if err != nil { + var fne *feature.FeatureNotFoundError + + if errors.As(err, &fne) || models.IsGenericNotFoundError(err) { + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureNotFound, + }, nil + } + + return governance.FeatureAccess{}, fmt.Errorf("failed to get feature %q: %w", featureKey, err) + } + + return governance.FeatureAccess{ + HasAccess: false, + Reason: governance.AccessReasonFeatureUnavailable, + }, nil +} diff --git a/openmeter/governance/service/service_test.go b/openmeter/governance/service/service_test.go new file mode 100644 index 0000000000..d19cf0fef2 --- /dev/null +++ b/openmeter/governance/service/service_test.go @@ -0,0 +1,622 @@ +package service + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/oklog/ulid/v2" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + metricnoop "go.opentelemetry.io/otel/metric/noop" + "go.opentelemetry.io/otel/trace/noop" + + "github.com/openmeterio/openmeter/app/config" + "github.com/openmeterio/openmeter/openmeter/customer" + customeradapter "github.com/openmeterio/openmeter/openmeter/customer/adapter" + customerservice "github.com/openmeterio/openmeter/openmeter/customer/service" + entdb "github.com/openmeterio/openmeter/openmeter/ent/db" + "github.com/openmeterio/openmeter/openmeter/entitlement" + "github.com/openmeterio/openmeter/openmeter/governance" + "github.com/openmeterio/openmeter/openmeter/meter" + meteradapter "github.com/openmeterio/openmeter/openmeter/meter/mockadapter" + "github.com/openmeterio/openmeter/openmeter/productcatalog/feature" + "github.com/openmeterio/openmeter/openmeter/registry" + registrybuilder "github.com/openmeterio/openmeter/openmeter/registry/builder" + streamingtestutils "github.com/openmeterio/openmeter/openmeter/streaming/testutils" + "github.com/openmeterio/openmeter/openmeter/subject" + subjectadapter "github.com/openmeterio/openmeter/openmeter/subject/adapter" + subjectservice "github.com/openmeterio/openmeter/openmeter/subject/service" + "github.com/openmeterio/openmeter/openmeter/testutils" + "github.com/openmeterio/openmeter/openmeter/watermill/eventbus" + "github.com/openmeterio/openmeter/pkg/clock" + "github.com/openmeterio/openmeter/pkg/datetime" + "github.com/openmeterio/openmeter/pkg/framework/lockr" + pagination "github.com/openmeterio/openmeter/pkg/pagination/v2" + "github.com/openmeterio/openmeter/pkg/timeutil" +) + +const testPageSize = 100 + +func newTestNamespace(t *testing.T) string { + t.Helper() + + return ulid.Make().String() +} + +// migrateOnce serializes schema migrations to avoid concurrent-write errors from ent. +var migrateOnce sync.Mutex + +type testDeps struct { + dbClient *testutils.TestDB + subjectService subject.Service + customerService customer.Service + meterService meter.ManageService + featureRepo feature.FeatureRepo + registry *registry.Entitlement + streamingConnector *streamingtestutils.MockStreamingConnector +} + +func (d *testDeps) close(t *testing.T) { + t.Helper() + + if err := d.dbClient.EntDriver.Close(); err != nil { + t.Errorf("close ent driver: %v", err) + } + + if err := d.dbClient.PGDriver.Close(); err != nil { + t.Errorf("close pg driver: %v", err) + } +} + +func setupTestDeps(t *testing.T) *testDeps { + t.Helper() + + logger := testutils.NewDiscardLogger(t) + testdb := testutils.InitPostgresDB(t) + dbClient := testdb.EntDriver.Client() + + migrateOnce.Lock() + require.NoError(t, dbClient.Schema.Create(context.Background())) + migrateOnce.Unlock() + + meterService, err := meteradapter.NewManage(nil) + require.NoError(t, err) + + subjectAdapter, err := subjectadapter.New(dbClient) + require.NoError(t, err) + + subjectSvc, err := subjectservice.New(subjectAdapter) + require.NoError(t, err) + + customerAdapter, err := customeradapter.New(customeradapter.Config{ + Client: dbClient, + Logger: logger, + }) + require.NoError(t, err) + + customerSvc, err := customerservice.New(customerservice.Config{ + Adapter: customerAdapter, + Publisher: eventbus.NewMock(t), + }) + require.NoError(t, err) + + locker, err := lockr.NewLocker(&lockr.LockerConfig{Logger: logger}) + require.NoError(t, err) + + streamingConnector := streamingtestutils.NewMockStreamingConnector(t) + + reg := registrybuilder.GetEntitlementRegistry(registrybuilder.EntitlementOptions{ + DatabaseClient: dbClient, + StreamingConnector: streamingConnector, + Logger: logger, + Tracer: noop.NewTracerProvider().Tracer("test"), + MeterService: meterService, + CustomerService: customerSvc, + Publisher: eventbus.NewMock(t), + EntitlementsConfiguration: config.EntitlementsConfiguration{ + GracePeriod: datetime.ISODurationString("P1D"), + }, + Locker: locker, + }) + + return &testDeps{ + dbClient: testdb, + subjectService: subjectSvc, + customerService: customerSvc, + meterService: meterService, + featureRepo: reg.FeatureRepo, + registry: reg, + streamingConnector: streamingConnector, + } +} + +func newTestService(t *testing.T, deps *testDeps) governance.Service { + t.Helper() + + svc, err := New(Config{ + Customer: deps.customerService, + Entitlement: deps.registry.Entitlement, + Feature: deps.registry.Feature, + Tracer: noop.NewTracerProvider().Tracer("test"), + Meter: metricnoop.NewMeterProvider().Meter("test"), + }) + require.NoError(t, err) + + return svc +} + +func createCustomer(t *testing.T, deps *testDeps, ns, key string, subjectKeys []string) *customer.Customer { + t.Helper() + + for _, sk := range subjectKeys { + _, err := deps.subjectService.Create(t.Context(), subject.CreateInput{ + Namespace: ns, + Key: sk, + }) + require.NoError(t, err) + } + + cust, err := deps.customerService.CreateCustomer(t.Context(), customer.CreateCustomerInput{ + Namespace: ns, + CustomerMutate: customer.CustomerMutate{ + Key: lo.ToPtr(key), + Name: key, + UsageAttribution: &customer.CustomerUsageAttribution{ + SubjectKeys: subjectKeys, + }, + }, + }) + require.NoError(t, err) + + return cust +} + +func createBooleanFeatureAndEntitlement(t *testing.T, deps *testDeps, ns, featureKey string, cust *customer.Customer) { + t.Helper() + + feat, err := deps.featureRepo.CreateFeature(t.Context(), feature.CreateFeatureInputs{ + Key: featureKey, + Name: featureKey, + Namespace: ns, + }) + require.NoError(t, err) + + _, err = deps.registry.Entitlement.CreateEntitlement(t.Context(), entitlement.CreateEntitlementInputs{ + Namespace: ns, + UsageAttribution: cust.GetUsageAttribution(), + FeatureKey: lo.ToPtr(featureKey), + FeatureID: lo.ToPtr(feat.ID), + EntitlementType: entitlement.EntitlementTypeBoolean, + }, nil) + require.NoError(t, err) +} + +func createOrphanFeature(t *testing.T, deps *testDeps, ns, featureKey string) { + t.Helper() + + _, err := deps.featureRepo.CreateFeature(t.Context(), feature.CreateFeatureInputs{ + Key: featureKey, + Name: featureKey, + Namespace: ns, + }) + require.NoError(t, err) +} + +// createMeterInPG writes a meter row to ent DB (FK constraint on features.meter_id). +// The mock meter adapter only stores in memory; this must be called after CreateMeter. +func createMeterInPG(t *testing.T, dbClient *entdb.Client, mtr meter.Meter) { + t.Helper() + + _, err := dbClient.Meter.Create(). + SetID(mtr.ID). + SetNamespace(mtr.Namespace). + SetName(mtr.Name). + SetKey(mtr.Key). + SetAggregation(mtr.Aggregation). + SetEventType(mtr.EventType). + SetNillableValueProperty(mtr.ValueProperty). + Save(t.Context()) + require.NoError(t, err) +} + +func createMeter(t *testing.T, deps *testDeps, ns, key string) meter.Meter { + t.Helper() + + mtr, err := deps.meterService.CreateMeter(t.Context(), meter.CreateMeterInput{ + Namespace: ns, + Name: key, + Key: key, + Aggregation: meter.MeterAggregationSum, + EventType: "test", + ValueProperty: lo.ToPtr("$.value"), + }) + require.NoError(t, err) + + createMeterInPG(t, deps.dbClient.EntDriver.Client(), mtr) + + return mtr +} + +func createMeteredFeatureAndEntitlement(t *testing.T, deps *testDeps, ns, featureKey string, mtr meter.Meter, cust *customer.Customer, issueAfterReset *float64) { + t.Helper() + + feat, err := deps.featureRepo.CreateFeature(t.Context(), feature.CreateFeatureInputs{ + Key: featureKey, + Name: featureKey, + Namespace: ns, + MeterID: lo.ToPtr(mtr.ID), + }) + require.NoError(t, err) + + _, err = deps.registry.Entitlement.CreateEntitlement(t.Context(), entitlement.CreateEntitlementInputs{ + Namespace: ns, + UsageAttribution: cust.GetUsageAttribution(), + FeatureKey: lo.ToPtr(featureKey), + FeatureID: lo.ToPtr(feat.ID), + EntitlementType: entitlement.EntitlementTypeMetered, + UsagePeriod: lo.ToPtr(entitlement.NewUsagePeriodInputFromRecurrence(timeutil.Recurrence{ + Interval: timeutil.RecurrencePeriodDaily, + Anchor: clock.Now(), + })), + IssueAfterReset: issueAfterReset, + }, nil) + require.NoError(t, err) +} + +// --- Tests --- + +func TestQueryAccess_UnknownCustomerKey(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"ghost"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + assert.Empty(t, res.Customers) + require.Len(t, res.Errors, 1) + assert.Equal(t, governance.QueryErrorCustomerNotFound, res.Errors[0].Code) + assert.Equal(t, "ghost", res.Errors[0].CustomerKey) +} + +func TestQueryAccess_KnownCustomerNoEntitlements(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{cust.GetUsageAttribution().SubjectKeys[0]}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + assert.Empty(t, res.Customers[0].Features) + assert.Empty(t, res.Errors) +} + +func TestQueryAccess_BooleanEntitlement_HasAccess(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + createBooleanFeatureAndEntitlement(t, deps, ns, "premium", cust) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + FeatureKeys: []string{"premium"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + assert.Empty(t, res.Errors) + + fa := res.Customers[0].Features["premium"] + + assert.True(t, fa.HasAccess) + assert.Nil(t, fa.Reason) +} + +func TestQueryAccess_FeatureNotFound(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + createCustomer(t, deps, ns, "acme", []string{"acme"}) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + FeatureKeys: []string{"does-not-exist"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + + fa := res.Customers[0].Features["does-not-exist"] + + assert.False(t, fa.HasAccess) + require.NotNil(t, fa.Reason) + assert.Equal(t, governance.ReasonCodeFeatureNotFound, fa.Reason.Code) +} + +func TestQueryAccess_FeatureUnavailable(t *testing.T) { + // Feature exists in org but customer has no entitlement for it. + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + createCustomer(t, deps, ns, "acme", []string{"acme"}) + createOrphanFeature(t, deps, ns, "enterprise") + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + FeatureKeys: []string{"enterprise"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + + fa := res.Customers[0].Features["enterprise"] + + assert.False(t, fa.HasAccess) + require.NotNil(t, fa.Reason) + assert.Equal(t, governance.ReasonCodeFeatureUnavailable, fa.Reason.Code) +} + +func TestQueryAccess_MultipleKeysSameCustomer(t *testing.T) { + // Two input keys resolve to the same customer; result has one entry with both keys in Matched. + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + // customer key = "acme", usage attribution subject key = "acme-sub" + createCustomer(t, deps, ns, "acme", []string{"acme-sub"}) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme", "acme-sub"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + assert.Empty(t, res.Errors) + require.Len(t, res.Customers, 1, "two keys resolving to same customer should collapse into one result") + assert.Len(t, res.Customers[0].Matched, 2) + assert.ElementsMatch(t, []string{"acme", "acme-sub"}, res.Customers[0].Matched) +} + +func TestQueryAccess_MixedHitsAndMisses(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + createBooleanFeatureAndEntitlement(t, deps, ns, "feature-a", cust) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme", "unknown-key"}, + FeatureKeys: []string{"feature-a"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + require.Len(t, res.Errors, 1) + assert.Equal(t, governance.QueryErrorCustomerNotFound, res.Errors[0].Code) + assert.True(t, res.Customers[0].Features["feature-a"].HasAccess) +} + +func TestQueryAccess_NoFeatureKeysReturnsAll(t *testing.T) { + // When no feature keys are given, all org features are returned — including ones + // the customer has no entitlement for (marked feature-unavailable). + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + createBooleanFeatureAndEntitlement(t, deps, ns, "feat-1", cust) + createBooleanFeatureAndEntitlement(t, deps, ns, "feat-2", cust) + // feat-3 exists in the org but the customer has no entitlement for it. + createOrphanFeature(t, deps, ns, "feat-3") + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + assert.Len(t, res.Customers[0].Features, 3) + assert.True(t, res.Customers[0].Features["feat-1"].HasAccess) + assert.True(t, res.Customers[0].Features["feat-2"].HasAccess) + + feat3 := res.Customers[0].Features["feat-3"] + + assert.False(t, feat3.HasAccess) + require.NotNil(t, feat3.Reason) + assert.Equal(t, governance.ReasonCodeFeatureUnavailable, feat3.Reason.Code) +} + +func TestQueryAccess_MeteredEntitlement_HasAccess(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + now := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + clock.SetTime(now) + defer clock.ResetTime() + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + mtr := createMeter(t, deps, ns, "api-calls") + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + // IssueAfterReset=10.0 → balance starts at 10, HasAccess=true + createMeteredFeatureAndEntitlement(t, deps, ns, "premium", mtr, cust, lo.ToPtr(10.0)) + + // Add an event so the streaming mock has data for the meter. + deps.streamingConnector.AddSimpleEvent(mtr.Key, 1, now) + + clock.SetTime(now.Add(time.Hour)) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + FeatureKeys: []string{"premium"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + assert.Empty(t, res.Errors) + + fa := res.Customers[0].Features["premium"] + + assert.True(t, fa.HasAccess) + assert.Nil(t, fa.Reason) +} + +func TestQueryAccess_MeteredEntitlement_Exhausted(t *testing.T) { + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + now := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + clock.SetTime(now) + defer clock.ResetTime() + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + mtr := createMeter(t, deps, ns, "api-calls") + cust := createCustomer(t, deps, ns, "acme", []string{"acme"}) + // No IssueAfterReset → balance=0, HasAccess=false → usage limit reached + createMeteredFeatureAndEntitlement(t, deps, ns, "premium", mtr, cust, nil) + + deps.streamingConnector.AddSimpleEvent(mtr.Key, 1, now) + + clock.SetTime(now.Add(time.Hour)) + + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: []string{"acme"}, + FeatureKeys: []string{"premium"}, + PageSize: testPageSize, + }) + require.NoError(t, err) + + require.Len(t, res.Customers, 1) + assert.Empty(t, res.Errors) + + fa := res.Customers[0].Features["premium"] + + assert.False(t, fa.HasAccess) + require.NotNil(t, fa.Reason) + assert.Equal(t, governance.ReasonCodeUsageLimitReached, fa.Reason.Code) +} + +func TestQueryAccess_Pagination(t *testing.T) { + // given: 3 customers (c1, c2, c3 in creation order); pageSize=1 + deps := setupTestDeps(t) + t.Cleanup(func() { deps.close(t) }) + + now := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC) + clock.SetTime(now) + defer clock.ResetTime() + + svc := newTestService(t, deps) + ns := newTestNamespace(t) + + createCustomer(t, deps, ns, "c1", []string{"c1"}) + clock.SetTime(now.Add(time.Second)) + createCustomer(t, deps, ns, "c2", []string{"c2"}) + clock.SetTime(now.Add(2 * time.Second)) + createCustomer(t, deps, ns, "c3", []string{"c3"}) + + allKeys := []string{"c1", "c2", "c3"} + + query := func(after, before *pagination.Cursor) governance.QueryResult { + res, err := svc.QueryAccess(t.Context(), governance.QueryAccessInput{ + Namespace: ns, + CustomerKeys: allKeys, + PageSize: 1, + After: after, + Before: before, + }) + require.NoError(t, err) + + return res + } + + // Page 1: [c1] — no previous, next set + page1 := query(nil, nil) + require.Len(t, page1.Customers, 1) + assert.Equal(t, "c1", page1.Customers[0].Matched[0]) + assert.False(t, page1.HasPrev, "no previous on first page") + require.True(t, page1.HasNext, "next must be set on page 1") + require.NotNil(t, page1.Last) + + // Page 2: [c2] — previous and next set. Forward uses the prior page's Last cursor. + page2 := query(page1.Last, nil) + require.Len(t, page2.Customers, 1) + assert.Equal(t, "c2", page2.Customers[0].Matched[0]) + require.True(t, page2.HasPrev, "previous must be set on page 2") + require.True(t, page2.HasNext, "next must be set on page 2") + require.NotNil(t, page2.First) + require.NotNil(t, page2.Last) + + // Page 3 (last): [c3] — previous set, no next + page3 := query(page2.Last, nil) + require.Len(t, page3.Customers, 1) + assert.Equal(t, "c3", page3.Customers[0].Matched[0]) + assert.True(t, page3.HasPrev, "previous must be set on last page") + assert.False(t, page3.HasNext, "no next on last page") + require.NotNil(t, page3.Last) + + // Cursor past end → empty page, no cursors + pastEnd := query(page3.Last, nil) + assert.Empty(t, pastEnd.Customers) + assert.False(t, pastEnd.HasNext) + assert.Nil(t, pastEnd.First) + assert.Nil(t, pastEnd.Last) + + // Backward from page 2's first (previous) cursor → [c1], no previous, next set + pageBack := query(nil, page2.First) + require.Len(t, pageBack.Customers, 1) + assert.Equal(t, "c1", pageBack.Customers[0].Matched[0]) + assert.False(t, pageBack.HasPrev, "no previous before c1") + assert.True(t, pageBack.HasNext, "next must be set in backward result") +} diff --git a/openmeter/server/router/router.go b/openmeter/server/router/router.go index a90c55d90b..426fb78e86 100644 --- a/openmeter/server/router/router.go +++ b/openmeter/server/router/router.go @@ -37,6 +37,7 @@ import ( entitlementdriver "github.com/openmeterio/openmeter/openmeter/entitlement/driver" entitlementdriverv2 "github.com/openmeterio/openmeter/openmeter/entitlement/driver/v2" meteredentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/metered" + "github.com/openmeterio/openmeter/openmeter/governance" infohttpdriver "github.com/openmeterio/openmeter/openmeter/info/httpdriver" "github.com/openmeterio/openmeter/openmeter/ingest" ingesthttpdriver "github.com/openmeterio/openmeter/openmeter/ingest/httpdriver" @@ -117,6 +118,7 @@ type Config struct { EntitlementBalanceConnector meteredentitlement.Connector ErrorHandler errorsx.Handler FeatureConnector feature.FeatureConnector + GovernanceService governance.Service GrantConnector credit.GrantConnector GrantRepo grant.Repo IngestService ingest.Service @@ -251,6 +253,10 @@ func (c Config) Validate() error { return errors.New("feature gate is required") } + if c.GovernanceService == nil { + return errors.New("governance service is required") + } + return nil } diff --git a/openmeter/server/server.go b/openmeter/server/server.go index c97ebde5e9..4fd650f670 100644 --- a/openmeter/server/server.go +++ b/openmeter/server/server.go @@ -131,6 +131,7 @@ func NewServer(config *Config) (*Server, error) { CustomerBalanceFacade: config.RouterConfig.CustomerBalanceFacade, CurrencyService: config.RouterConfig.CurrencyService, EntitlementService: config.RouterConfig.EntitlementConnector, + GovernanceService: config.RouterConfig.GovernanceService, IngestService: config.RouterConfig.IngestService, MeterEventService: config.RouterConfig.MeterEventService, LLMCostService: config.RouterConfig.LLMCostService, diff --git a/openmeter/server/server_test.go b/openmeter/server/server_test.go index 4b6d828c3c..05bfc46bcf 100644 --- a/openmeter/server/server_test.go +++ b/openmeter/server/server_test.go @@ -17,6 +17,8 @@ import ( "github.com/samber/lo" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + metricnoop "go.opentelemetry.io/otel/metric/noop" + "go.opentelemetry.io/otel/trace/noop" "github.com/openmeterio/openmeter/api" apiv3 "github.com/openmeterio/openmeter/api/v3" @@ -34,6 +36,7 @@ import ( "github.com/openmeterio/openmeter/openmeter/customer" "github.com/openmeterio/openmeter/openmeter/entitlement" meteredentitlement "github.com/openmeterio/openmeter/openmeter/entitlement/metered" + governanceservice "github.com/openmeterio/openmeter/openmeter/governance/service" "github.com/openmeterio/openmeter/openmeter/ingest" "github.com/openmeterio/openmeter/openmeter/llmcost" "github.com/openmeterio/openmeter/openmeter/meter" @@ -743,6 +746,16 @@ func getTestServer(t *testing.T, opts ...func(*router.Config)) (*Server, *MockSt // Create subject service subjectService := &NoopSubjectService{} + // Create governance service from the noop collaborators + governanceService, err := governanceservice.New(governanceservice.Config{ + Customer: customerService, + Entitlement: &NoopEntitlementConnector{}, + Feature: featureService, + Tracer: noop.NewTracerProvider().Tracer("test"), + Meter: metricnoop.NewMeterProvider().Meter("test"), + }) + assert.NoError(t, err, "failed to create governance service") + config := &Config{ RouterConfig: router.Config{ Addon: addonService, @@ -758,6 +771,7 @@ func getTestServer(t *testing.T, opts ...func(*router.Config)) (*Server, *MockSt EntitlementBalanceConnector: &NoopEntitlementBalanceConnector{}, ErrorHandler: errorsx.NopHandler{}, FeatureConnector: featureService, + GovernanceService: governanceService, GrantConnector: &NoopGrantConnector{}, // Use the grant repo GrantRepo: grantRepo,