feat(charges): ubp flatfee costbasis support

Author: turip

app/common/billing.go

@@ -173,6 +173,7 @@ func NewBillingRegistry(
 	tracer trace.Tracer,
 	taxCodeService taxcode.Service,
 	currencyResolver currencies.CurrencyResolver,
+	currenciesService currencies.Service,
 	locker *lockr.Locker,
 	ledgerService ledger.Ledger,
 	balanceQuerier ledger.BalanceQuerier,
@@ -226,6 +227,7 @@ func NewBillingRegistry(
 			breakageService,
 			taxCodeService,
 			currencyResolver,
+			currenciesService,
 			fsConfig.NamespaceLockdown,
 			creditsConfig,
 			featureGate,

app/common/charges.go

@@ -212,6 +212,7 @@ func NewChargesFlatFeeService(
 	metaAdapter meta.Adapter,
 	locker *lockr.Locker,
 	ratingService rating.Service,
+	currenciesService currencies.Service,
 ) (flatfee.Service, error) {
 	flatFeeSvc, err := flatfeeservice.New(flatfeeservice.Config{
 		Adapter:       flatFeeAdapter,
@@ -220,6 +221,7 @@ func NewChargesFlatFeeService(
 		MetaAdapter:   metaAdapter,
 		Locker:        locker,
 		RatingService: ratingService,
+		Currencies:    currenciesService,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to create charges flat fee service: %w", err)
@@ -414,6 +416,7 @@ func newChargesRegistry(
 	breakageService ledgerbreakage.Service,
 	taxCodeService taxcode.Service,
 	currencyResolver currencies.CurrencyResolver,
+	currenciesService currencies.Service,
 	fsNamespaceLockdown []string,
 	creditsConfig config.CreditsConfiguration,
 	featureGate *featuregate.FeatureGateChecker,
@@ -473,7 +476,7 @@ func newChargesRegistry(
 		return nil, err
 	}
 
-	flatFeeSvc, err := NewChargesFlatFeeService(flatFeeAdapter, flatFeeHandler, lineageService, metaAdapter, locker, ratingService)
+	flatFeeSvc, err := NewChargesFlatFeeService(flatFeeAdapter, flatFeeHandler, lineageService, metaAdapter, locker, ratingService, currenciesService)
 	if err != nil {
 		return nil, err
 	}

cmd/billing-worker/wire_gen.go

@@ -9,6 +9,7 @@ import (
 	"github.com/alpacahq/alpacadecimal"
 
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
+	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/costbasis"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/creditrealization"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/invoicedusage"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/payment"
@@ -24,10 +25,15 @@ type Adapter interface {
 	ChargeRunAdapter
 	ChargeInvoicedUsageAdapter
 	ChargePaymentAdapter
+	ChargeCostBasisAdapter
 
 	entutils.TxCreator
 }
 
+type ChargeCostBasisAdapter interface {
+	SetResolvedCostBasis(ctx context.Context, input costbasis.SetResolvedCostBasisInput) (costbasis.CostBasis, error)
+}
+
 type ChargeAdapter interface {
 	CreateCharges(ctx context.Context, charges CreateChargesInput) ([]Charge, error)
 
@@ -141,6 +147,7 @@ type IntentWithInitialStatus struct {
 	InitialAdvanceAfter       *time.Time
 	AmountAfterProration      alpacadecimal.Decimal
 	NoFiatTransactionRequired bool
+	ResolvedCostBasis         *costbasis.State
 }
 
 func (i IntentWithInitialStatus) Validate() error {

cmd/jobs/internal/wire_gen.go

@@ -5,13 +5,15 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/oklog/ulid/v2"
 	"github.com/samber/lo"
 
 	"github.com/openmeterio/openmeter/openmeter/billing"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/flatfee"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/meta"
 	metaadapter "github.com/openmeterio/openmeter/openmeter/billing/charges/meta/adapter"
 	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/chargemeta"
+	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/costbasis"
 	"github.com/openmeterio/openmeter/openmeter/ent/db"
 	dbchargeflatfee "github.com/openmeterio/openmeter/openmeter/ent/db/chargeflatfee"
 	dbchargeflatfeeoverride "github.com/openmeterio/openmeter/openmeter/ent/db/chargeflatfeeoverride"
@@ -86,6 +88,10 @@ func (a *adapter) UpdateCharge(ctx context.Context, charge flatfee.ChargeBase) (
 			return flatfee.ChargeBase{}, err
 		}
 
+		if err := tx.loadCostBasisEdge(ctx, dbUpdatedChargeBase); err != nil {
+			return flatfee.ChargeBase{}, err
+		}
+
 		if overrideLayer := charge.Intent.GetOverrideLayerMutableFields(); overrideLayer != nil {
 			intentOverride, err := tx.updateIntentOverride(ctx, charge.GetChargeID(), overrideLayer, intent.Currency)
 			if err != nil {
@@ -132,6 +138,9 @@ func (a *adapter) UpdateSubscriptionItemID(ctx context.Context, charge flatfee.C
 		}
 
 		updatedChargeBase.Edges.IntentOverride = override
+		if err := tx.loadCostBasisEdge(ctx, updatedChargeBase); err != nil {
+			return flatfee.Charge{}, err
+		}
 		mappedChargeBase, err := fromDBBaseWithCurrency(updatedChargeBase, charge.Intent.GetBaseIntent().Currency)
 		if err != nil {
 			return flatfee.Charge{}, err
@@ -208,19 +217,69 @@ func (a *adapter) CreateCharges(ctx context.Context, in flatfee.CreateChargesInp
 	}
 
 	return entutils.TransactingRepo(ctx, a, func(ctx context.Context, tx *adapter) ([]flatfee.Charge, error) {
-		creates, err := slicesx.MapWithErr(in.Intents, func(intent flatfee.IntentWithInitialStatus) (*db.ChargeFlatFeeCreate, error) {
-			return tx.buildCreateFlatFeeCharge(in.Namespace, intent)
+		type preparedCreate struct {
+			costBasis *db.ChargeFlatFeeCostBasisCreate
+			charge    *db.ChargeFlatFeeCreate
+		}
+
+		preparedCreates := make([]preparedCreate, 0, len(in.Intents))
+		for _, intent := range in.Intents {
+			chargeCreate, err := tx.buildCreateFlatFeeCharge(in.Namespace, intent)
+			if err != nil {
+				return nil, err
+			}
+
+			var costBasisCreate *db.ChargeFlatFeeCostBasisCreate
+			if intent.Intent.CostBasis != nil {
+				costBasisCreate, err = costbasis.Create(tx.db.ChargeFlatFeeCostBasis.Create(), costbasis.CreateInput{
+					NamespacedID: models.NamespacedID{
+						Namespace: in.Namespace,
+						ID:        ulid.Make().String(),
+					},
+					CurrencyID: intent.Intent.Currency.ID,
+					Intent:     *intent.Intent.CostBasis,
+					State:      intent.ResolvedCostBasis,
+				})
+				if err != nil {
+					return nil, fmt.Errorf("building flat fee cost basis: %w", err)
+				}
+			}
+
+			preparedCreates = append(preparedCreates, preparedCreate{
+				costBasis: costBasisCreate,
+				charge:    chargeCreate,
+			})
+		}
+
+		costBasisCreates := lo.Filter(preparedCreates, func(create preparedCreate, _ int) bool {
+			return create.costBasis != nil
 		})
-		if err != nil {
-			return nil, err
+
+		var createdCostBases []*db.ChargeFlatFeeCostBasis
+		if len(costBasisCreates) > 0 {
+			var err error
+			createdCostBases, err = tx.db.ChargeFlatFeeCostBasis.CreateBulk(
+				lo.Map(costBasisCreates, func(create preparedCreate, _ int) *db.ChargeFlatFeeCostBasisCreate {
+					return create.costBasis
+				})...,
+			).Save(ctx)
+			if err != nil {
+				return nil, fmt.Errorf("creating flat fee cost bases: %w", err)
+			}
+
+			lo.ForEach(costBasisCreates, func(create preparedCreate, idx int) {
+				create.charge.SetCostBasisID(createdCostBases[idx].ID)
+			})
 		}
 
-		entities, err := tx.db.ChargeFlatFee.CreateBulk(creates...).Save(ctx)
+		chargeCreates := lo.Map(preparedCreates, func(create preparedCreate, _ int) *db.ChargeFlatFeeCreate {
+			return create.charge
+		})
+		entities, err := tx.db.ChargeFlatFee.CreateBulk(chargeCreates...).Save(ctx)
 		if err != nil {
 			return nil, metaadapter.MapChargeConstraintError(err)
 		}
 
-		// Let's reserve the charge IDs
 		err = tx.metaAdapter.RegisterCharges(ctx, meta.RegisterChargesInput{
 			Namespace: in.Namespace,
 			Type:      meta.ChargeTypeFlatFee,
@@ -235,7 +294,20 @@ func (a *adapter) CreateCharges(ctx context.Context, in flatfee.CreateChargesInp
 			return nil, err
 		}
 
+		costBasisByID := lo.SliceToMap(createdCostBases, func(entity *db.ChargeFlatFeeCostBasis) (string, *db.ChargeFlatFeeCostBasis) {
+			return entity.ID, entity
+		})
+
 		return lo.MapErr(entities, func(entity *db.ChargeFlatFee, idx int) (flatfee.Charge, error) {
+			if entity.CostBasisID != nil {
+				createdCostBasis, ok := costBasisByID[*entity.CostBasisID]
+				if !ok {
+					return flatfee.Charge{}, fmt.Errorf("created flat fee cost basis %s not found", *entity.CostBasisID)
+				}
+
+				entity.Edges.CostBasis = createdCostBasis
+			}
+
 			return FromDBWithCurrency(entity, in.Intents[idx].Intent.Currency, meta.ExpandNone)
 		})
 	})
@@ -251,7 +323,8 @@ func (a *adapter) GetByIDs(ctx context.Context, input flatfee.GetByIDsInput) ([]
 			Where(dbchargeflatfee.Namespace(input.Namespace)).
 			Where(dbchargeflatfee.IDIn(input.IDs...)).
 			WithIntentOverride().
-			WithCustomCurrency()
+			WithCustomCurrency().
+			WithCostBasis()
 
 		if input.Expands.Has(meta.ExpandRealizations) {
 			query = expandRealizations(query)
@@ -294,7 +367,8 @@ func (a *adapter) GetByID(ctx context.Context, input flatfee.GetByIDInput) (flat
 			Where(dbchargeflatfee.Namespace(input.ChargeID.Namespace)).
 			Where(dbchargeflatfee.ID(input.ChargeID.ID)).
 			WithIntentOverride().
-			WithCustomCurrency()
+			WithCustomCurrency().
+			WithCostBasis()
 
 		if input.Expands.Has(meta.ExpandRealizations) {
 			query = expandRealizations(query)

cmd/server/wire_gen.go

@@ -0,0 +1,97 @@
+package adapter
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/openmeterio/openmeter/openmeter/billing/charges/flatfee"
+	"github.com/openmeterio/openmeter/openmeter/billing/charges/models/costbasis"
+	entdb "github.com/openmeterio/openmeter/openmeter/ent/db"
+	dbchargeflatfeecostbasis "github.com/openmeterio/openmeter/openmeter/ent/db/chargeflatfeecostbasis"
+	"github.com/openmeterio/openmeter/pkg/framework/entutils"
+	"github.com/openmeterio/openmeter/pkg/models"
+)
+
+var _ flatfee.ChargeCostBasisAdapter = (*adapter)(nil)
+
+func (a *adapter) loadCostBasisEdge(ctx context.Context, entity *entdb.ChargeFlatFee) error {
+	if entity.Edges.CostBasis != nil {
+		return fmt.Errorf("flat fee cost basis edge is already loaded [charge_id=%s,edge_id=%s]", entity.ID, entity.Edges.CostBasis.ID)
+	}
+
+	if entity.CostBasisID == nil {
+		return nil
+	}
+
+	costBasisEntity, err := a.db.ChargeFlatFeeCostBasis.Query().
+		Where(
+			dbchargeflatfeecostbasis.ID(*entity.CostBasisID),
+			dbchargeflatfeecostbasis.Namespace(entity.Namespace),
+		).
+		Only(ctx)
+	if err != nil {
+		return fmt.Errorf("load flat fee cost basis edge [charge_id=%s,cost_basis_id=%s]: %w", entity.ID, *entity.CostBasisID, err)
+	}
+
+	entity.Edges.CostBasis = costBasisEntity
+
+	return nil
+}
+
+func (a *adapter) SetResolvedCostBasis(ctx context.Context, input costbasis.SetResolvedCostBasisInput) (costbasis.CostBasis, error) {
+	if err := input.Validate(); err != nil {
+		return costbasis.CostBasis{}, err
+	}
+
+	return entutils.TransactingRepo(ctx, a, func(ctx context.Context, tx *adapter) (costbasis.CostBasis, error) {
+		update := tx.db.ChargeFlatFeeCostBasis.UpdateOneID(input.ID).
+			Where(
+				dbchargeflatfeecostbasis.Namespace(input.Namespace),
+				dbchargeflatfeecostbasis.ModeEQ(costbasis.ModeDynamic),
+				dbchargeflatfeecostbasis.ResolvedCostBasisIDIsNil(),
+				dbchargeflatfeecostbasis.ResolvedCostBasisIsNil(),
+				dbchargeflatfeecostbasis.ResolvedAtIsNil(),
+			)
+
+		update, err := costbasis.Set(update, input.State)
+		if err != nil {
+			return costbasis.CostBasis{}, err
+		}
+
+		entity, err := update.Save(ctx)
+		if entdb.IsNotFound(err) {
+			existingEntity, getErr := tx.db.ChargeFlatFeeCostBasis.Query().
+				Where(
+					dbchargeflatfeecostbasis.ID(input.ID),
+					dbchargeflatfeecostbasis.Namespace(input.Namespace),
+				).
+				Only(ctx)
+			if entdb.IsNotFound(getErr) {
+				return costbasis.CostBasis{}, models.NewGenericNotFoundError(
+					fmt.Errorf("flat fee cost basis not found: %s", input.ID),
+				)
+			}
+			if getErr != nil {
+				return costbasis.CostBasis{}, fmt.Errorf("get flat fee cost basis after conditional resolution: %w", getErr)
+			}
+
+			existing, getErr := costbasis.Get(existingEntity)
+			if getErr != nil {
+				return costbasis.CostBasis{}, fmt.Errorf("map flat fee cost basis after conditional resolution: %w", getErr)
+			}
+
+			if existing.Intent.Kind() == costbasis.ModeDynamic && existing.State != nil {
+				return existing, nil
+			}
+
+			return costbasis.CostBasis{}, models.NewGenericValidationError(
+				fmt.Errorf("flat fee cost basis is not an unresolved dynamic cost basis: %s", input.ID),
+			)
+		}
+		if err != nil {
+			return costbasis.CostBasis{}, fmt.Errorf("set resolved flat fee cost basis: %w", err)
+		}
+
+		return costbasis.Get(entity)
+	})
+}

openmeter/billing/charges/flatfee/adapter.go

@@ -83,6 +83,9 @@ func (a *adapter) CreateChargeOverride(ctx context.Context, charge flatfee.Charg
 		}
 
 		dbCharge.Edges.IntentOverride = dbIntentOverride
+		if err := tx.loadCostBasisEdge(ctx, dbCharge); err != nil {
+			return flatfee.ChargeBase{}, err
+		}
 
 		return fromDBBaseWithCurrency(dbCharge, charge.Intent.GetBaseIntent().Currency)
 	})