feat: implement realtime expand (#4155)

Author: turip

api/v3/handlers/customers/charges/convert.go

@@ -157,9 +157,15 @@ func convertChargeToAPI(charge billingcharges.Charge) (api.BillingCharge, error)
 
 // convertUsageBasedChargeTotals aggregates booked totals from persisted realization runs.
 func convertUsageBasedChargeTotals(charge usagebased.Charge) api.BillingChargeTotals {
-	return api.BillingChargeTotals{
+	out := api.BillingChargeTotals{
 		Booked: toAPIBillingTotals(charge.Realizations.Sum()),
 	}
+
+	if charge.Expands.RealtimeUsage != nil {
+		out.Realtime = lo.ToPtr(toAPIBillingTotals(*charge.Expands.RealtimeUsage))
+	}
+
+	return out
 }
 
 // toAPIBillingTotals maps a domain totals.Totals to the API BillingTotals type.

api/v3/handlers/customers/charges/list.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"slices"
 	"strings"
 
 	"github.com/samber/lo"
@@ -56,14 +57,19 @@ func (h *handler) ListCustomerCharges() ListCustomerChargesHandler {
 				})
 			}
 
+			// Realization runs are always required to compute booked totals.
+			expands := meta.Expands{meta.ExpandRealizations}
+			if args.Params.Expand != nil && slices.Contains(*args.Params.Expand, api.BillingChargesExpandRealTimeUsage) {
+				expands = expands.With(meta.ExpandRealtimeUsage)
+			}
+
 			req := ListCustomerChargesRequest{
 				Page:        page,
 				Namespace:   ns,
 				CustomerIDs: []string{args.CustomerID},
 				// Credit purchases are served by the credit grants API; exclude them here.
 				ChargeTypes: []meta.ChargeType{meta.ChargeTypeFlatFee, meta.ChargeTypeUsageBased},
-				// Realization runs are always required to compute booked totals.
-				Expands: meta.Expands{meta.ExpandRealizations},
+				Expands:     expands,
 			}
 
 			// Parse sort. When omitted, the service defaults to created_at ascending

openmeter/billing/charges/adapter/search.go

@@ -70,7 +70,9 @@ func (a *adapter) ListCharges(ctx context.Context, input charges.ListChargesInpu
 
 		if len(input.StatusIn) > 0 {
 			query = query.Where(dbchargessearchv1.StatusIn(input.StatusIn...))
-		} else if len(input.StatusNotIn) > 0 {
+		}
+
+		if len(input.StatusNotIn) > 0 {
 			query = query.Where(dbchargessearchv1.StatusNotIn(input.StatusNotIn...))
 		}
 

openmeter/billing/charges/meta/charge.go

@@ -67,12 +67,14 @@ func (t ChargeType) Validate() error {
 type Expand string
 
 const (
-	ExpandRealizations Expand = "realizations"
+	ExpandRealizations  Expand = "realizations"
+	ExpandRealtimeUsage Expand = "realtime_usage"
 )
 
 func (e Expand) Values() []Expand {
 	return []Expand{
 		ExpandRealizations,
+		ExpandRealtimeUsage,
 	}
 }
 

openmeter/billing/charges/service.go

@@ -144,13 +144,9 @@ type ListChargesInput struct {
 	CustomerIDs     []string
 	SubscriptionIDs []string
 	ChargeTypes     []meta.ChargeType
-	// StatusIn filters to only charges with one of the given statuses.
-	// Takes precedence over StatusNotIn when both are set.
-	// When empty and StatusNotIn is also empty, deleted charges are still
-	// excluded via the IncludeDeleted flag.
-	StatusIn       []meta.ChargeStatus
-	StatusNotIn    []meta.ChargeStatus
-	IncludeDeleted bool
+	StatusIn        []meta.ChargeStatus
+	StatusNotIn     []meta.ChargeStatus
+	IncludeDeleted  bool
 
 	// OrderBy is the field to sort by. Supported values: id, created_at,
 	// service_period.from, billing_period.from.
@@ -198,6 +194,10 @@ func (i ListChargesInput) Validate() error {
 		}
 	}
 
+	if len(i.StatusIn) > 0 && len(i.StatusNotIn) > 0 {
+		errs = append(errs, errors.New("status_in and status_not_in cannot be set at the same time"))
+	}
+
 	if err := i.Expands.Validate(); err != nil {
 		errs = append(errs, fmt.Errorf("expands: %w", err))
 	}

openmeter/billing/charges/usagebased/charge.go

@@ -8,6 +8,8 @@ import (
 	"github.com/samber/lo"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
+	"github.com/openmeterio/openmeter/openmeter/billing/models/totals"
+	"github.com/openmeterio/openmeter/openmeter/customer"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog"
 	"github.com/openmeterio/openmeter/openmeter/productcatalog/feature"
 	"github.com/openmeterio/openmeter/pkg/models"
@@ -68,6 +70,7 @@ type Charge struct {
 	ChargeBase
 
 	Realizations RealizationRuns `json:"realizations"`
+	Expands      Expands         `json:"expands"`
 }
 
 type Charges []Charge
@@ -79,6 +82,10 @@ func (c Charge) Validate() error {
 		errs = append(errs, fmt.Errorf("charge base: %w", err))
 	}
 
+	if err := c.Expands.Validate(); err != nil {
+		errs = append(errs, fmt.Errorf("expands: %w", err))
+	}
+
 	return models.NewNillableGenericValidationError(errors.Join(errs...))
 }
 
@@ -90,6 +97,13 @@ func (c Charge) GetCurrentRealizationRun() (RealizationRun, error) {
 	return c.Realizations.GetByID(*c.State.CurrentRealizationRunID)
 }
 
+func (c Charge) GetCustomerID() customer.CustomerID {
+	return customer.CustomerID{
+		Namespace: c.Namespace,
+		ID:        c.Intent.CustomerID,
+	}
+}
+
 func (c Charge) GetFeatureKeyOrID() ref.IDOrKey {
 	switch c.Status {
 	case StatusCreated:
@@ -213,3 +227,19 @@ func (s State) Validate() error {
 
 	return models.NewNillableGenericValidationError(errors.Join(errs...))
 }
+
+type Expands struct {
+	RealtimeUsage *totals.Totals `json:"realtimeUsage,omitempty"`
+}
+
+func (e Expands) Validate() error {
+	var errs []error
+
+	if e.RealtimeUsage != nil {
+		if err := e.RealtimeUsage.Validate(); err != nil {
+			errs = append(errs, fmt.Errorf("realtime usage: %w", err))
+		}
+	}
+
+	return models.NewNillableGenericValidationError(errors.Join(errs...))
+}

openmeter/billing/charges/usagebased/service/get.go

@@ -2,9 +2,27 @@ package service
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"sync"
 
+	"github.com/samber/lo"
+	"golang.org/x/sync/semaphore"
+
+	"github.com/openmeterio/openmeter/openmeter/billing"
+	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/usagebased"
+	usagebasedrating "github.com/openmeterio/openmeter/openmeter/billing/charges/usagebased/service/rating"
+	"github.com/openmeterio/openmeter/openmeter/customer"
+	"github.com/openmeterio/openmeter/pkg/clock"
 	"github.com/openmeterio/openmeter/pkg/framework/transaction"
+	"github.com/openmeterio/openmeter/pkg/ref"
+	"github.com/openmeterio/openmeter/pkg/slicesx"
+)
+
+const (
+	// defaultMaxParallelRatingsPerRequest is the number of workers to use for the rating (fetching from CH).
+	defaultMaxParallelRatingsPerRequest = 5
 )
 
 func (s *service) GetByIDs(ctx context.Context, input usagebased.GetByIDsInput) ([]usagebased.Charge, error) {
@@ -13,7 +31,19 @@ func (s *service) GetByIDs(ctx context.Context, input usagebased.GetByIDsInput)
 	}
 
 	return transaction.Run(ctx, s.adapter, func(ctx context.Context) ([]usagebased.Charge, error) {
-		return s.adapter.GetByIDs(ctx, input)
+		charges, err := s.adapter.GetByIDs(ctx, input)
+		if err != nil {
+			return nil, err
+		}
+
+		if input.Expands.Has(meta.ExpandRealtimeUsage) {
+			charges, err = s.expandChargesUsage(ctx, input.Namespace, charges)
+			if err != nil {
+				return nil, err
+			}
+		}
+
+		return charges, nil
 	})
 }
 
@@ -23,6 +53,138 @@ func (s *service) GetByID(ctx context.Context, input usagebased.GetByIDInput) (u
 	}
 
 	return transaction.Run(ctx, s.adapter, func(ctx context.Context) (usagebased.Charge, error) {
-		return s.adapter.GetByID(ctx, input)
+		charge, err := s.adapter.GetByID(ctx, input)
+		if err != nil {
+			return usagebased.Charge{}, err
+		}
+
+		if input.Expands.Has(meta.ExpandRealtimeUsage) {
+			totals, err := s.GetCurrentTotals(ctx, usagebased.GetCurrentTotalsInput{
+				ChargeID: charge.GetChargeID(),
+			})
+			if err != nil {
+				return usagebased.Charge{}, err
+			}
+
+			charge.Expands.RealtimeUsage = &totals.DueTotals
+		}
+
+		return charge, nil
+	})
+}
+
+func (s *service) expandChargesUsage(ctx context.Context, namespace string, charges usagebased.Charges) (usagebased.Charges, error) {
+	// Fetch unique customers from the charges to avoid duplicate calls to the customer override service.
+	uniqueCustomers := lo.Uniq(lo.Map(charges, func(charge usagebased.Charge, _ int) customer.CustomerID {
+		return charge.GetCustomerID()
+	}))
+
+	customerOverridesById := make(map[customer.CustomerID]billing.CustomerOverrideWithDetails)
+	for _, customerID := range uniqueCustomers {
+		customerOverride, err := s.customerOverrideService.GetCustomerOverride(ctx, billing.GetCustomerOverrideInput{
+			Customer: customerID,
+			Expand: billing.CustomerOverrideExpand{
+				Customer: true,
+			},
+		})
+		if err != nil {
+			return nil, err
+		}
+		customerOverridesById[customerID] = customerOverride
+	}
+
+	// Fetch all references featureMeters in bulk
+	referencedFeatureMeters := lo.Uniq(lo.Map(charges, func(charge usagebased.Charge, _ int) ref.IDOrKey {
+		return charge.GetFeatureKeyOrID()
+	}))
+
+	featureMeters, err := s.featureService.ResolveFeatureMeters(ctx, namespace, referencedFeatureMeters...)
+	if err != nil {
+		return nil, err
+	}
+
+	// Let's do the rating for each charge
+	sem := semaphore.NewWeighted(int64(defaultMaxParallelRatingsPerRequest))
+	storedAt := clock.Now()
+
+	errCh := make(chan error, len(charges))
+	ratingResults := sync.Map{}
+
+	var wg sync.WaitGroup
+
+	for _, charge := range charges {
+		featureMeter, err := charge.ResolveFeatureMeter(featureMeters)
+		if err != nil {
+			errCh <- fmt.Errorf("resolving feature meter: %w", err)
+			break
+		}
+
+		err = sem.Acquire(ctx, 1)
+		if err != nil {
+			// Clean up and stop the loop
+			errCh <- fmt.Errorf("acquiring worker slot: %w", err)
+			break
+		}
+
+		wg.Go(func() {
+			defer sem.Release(1)
+			var err error
+			defer func() {
+				if err != nil {
+					errCh <- err
+				}
+			}()
+
+			defer func() {
+				if r := recover(); r != nil {
+					err = fmt.Errorf("rating charge %s: %v", charge.ID, r)
+				}
+			}()
+
+			var ratingResult usagebasedrating.GetRatingForUsageResult
+			ratingResult, err = s.rater.GetRatingForUsage(ctx, usagebasedrating.GetRatingForUsageInput{
+				Charge:         charge,
+				Customer:       customerOverridesById[charge.GetCustomerID()],
+				FeatureMeter:   featureMeter,
+				StoredAtOffset: storedAt,
+			})
+			if err != nil {
+				err = fmt.Errorf("rating charge %s: %w", charge.ID, err)
+				return
+			}
+
+			ratingResults.Store(charge.GetChargeID(), ratingResult)
+		})
+	}
+
+	wg.Wait()
+
+	close(errCh)
+
+	var errs []error
+
+	for err := range errCh {
+		if err != nil {
+			errs = append(errs, err)
+		}
+	}
+
+	if len(errs) > 0 {
+		return nil, errors.Join(errs...)
+	}
+
+	return slicesx.MapWithErr(charges, func(charge usagebased.Charge) (usagebased.Charge, error) {
+		ratingResultAny, ok := ratingResults.Load(charge.GetChargeID())
+		if !ok {
+			return charge, fmt.Errorf("rating result not found for charge %s", charge.ID)
+		}
+
+		ratingResult, ok := ratingResultAny.(usagebasedrating.GetRatingForUsageResult)
+		if !ok {
+			return charge, fmt.Errorf("rating result not found for charge %s", charge.ID)
+		}
+
+		charge.Expands.RealtimeUsage = &ratingResult.Totals
+		return charge, nil
 	})
 }