feat: persist invoice id for ubp runs (#4306)

Author: turip

openmeter/billing/charges/service/usagebased_test.go

@@ -208,6 +208,8 @@ func (s *UsageBasedChargesTestSuite) TestUsageBasedCreditThenInvoicePartialInvoi
 		s.Equal(usagebased.RealizationRunTypePartialInvoice, currentRun.Type)
 		s.Require().NotNil(currentRun.LineID)
 		s.Equal(stdLine.ID, *currentRun.LineID)
+		s.Require().NotNil(currentRun.InvoiceID)
+		s.Equal(partialInvoice.ID, *currentRun.InvoiceID)
 		s.True(midPeriodInvoiceAt.Equal(currentRun.ServicePeriodTo))
 		s.True(midPeriodInvoiceAt.Equal(currentRun.StoredAtLT))
 		s.Require().NotNil(partialInvoice.CollectionAt)
@@ -335,6 +337,8 @@ func (s *UsageBasedChargesTestSuite) TestUsageBasedCreditThenInvoicePartialInvoi
 		currentRun, runErr := charge.GetCurrentRealizationRun()
 		s.NoError(runErr)
 		s.Equal(usagebased.RealizationRunTypeFinalRealization, currentRun.Type)
+		s.Require().NotNil(currentRun.InvoiceID)
+		s.Equal(finalInvoice.ID, *currentRun.InvoiceID)
 
 		// given
 		clock.FreezeTime(finalInvoice.DefaultCollectionAtForStandardInvoice())

openmeter/billing/charges/usagebased/adapter/mapper.go

@@ -98,6 +98,7 @@ func MapRealizationRunBaseFromDB(dbRun *entdb.ChargeUsageBasedRuns) usagebased.R
 
 		FeatureID:                 dbRun.FeatureID,
 		LineID:                    dbRun.LineID,
+		InvoiceID:                 dbRun.InvoiceID,
 		Type:                      dbRun.Type,
 		StoredAtLT:                dbRun.StoredAtLt.UTC(),
 		ServicePeriodTo:           dbRun.ServicePeriodTo.UTC(),

openmeter/billing/charges/usagebased/adapter/realizationrun.go

@@ -31,6 +31,7 @@ func (a *adapter) CreateRealizationRun(ctx context.Context, chargeID meta.Charge
 			SetServicePeriodTo(meta.NormalizeTimestamp(input.ServicePeriodTo)).
 			SetDetailedLinesPresent(false).
 			SetNillableBillingInvoiceLineID(input.LineID).
+			SetNillableBillingInvoiceID(input.InvoiceID).
 			SetMeteredQuantity(input.MeteredQuantity).
 			SetNoFiatTransactionRequired(input.NoFiatTransactionRequired)
 

openmeter/billing/charges/usagebased/realizationrun.go

@@ -65,6 +65,7 @@ type CreateRealizationRunInput struct {
 	StoredAtLT                time.Time             `json:"storedAtLT"`
 	ServicePeriodTo           time.Time             `json:"servicePeriodTo"`
 	LineID                    *string               `json:"lineId,omitempty"`
+	InvoiceID                 *string               `json:"invoiceId,omitempty"`
 	MeteredQuantity           alpacadecimal.Decimal `json:"meteredQuantity"`
 	Totals                    totals.Totals         `json:"totals"`
 	NoFiatTransactionRequired bool                  `json:"noFiatTransactionRequired"`
@@ -108,6 +109,10 @@ func (r CreateRealizationRunInput) Validate() error {
 		errs = append(errs, fmt.Errorf("line id must be non-empty"))
 	}
 
+	if r.InvoiceID != nil && *r.InvoiceID == "" {
+		errs = append(errs, fmt.Errorf("invoice id must be non-empty"))
+	}
+
 	return models.NewNillableGenericValidationError(errors.Join(errs...))
 }
 
@@ -167,6 +172,7 @@ type RealizationRunBase struct {
 
 	FeatureID string  `json:"featureId"`
 	LineID    *string `json:"lineId,omitempty"`
+	InvoiceID *string `json:"invoiceId,omitempty"`
 
 	Type       RealizationRunType `json:"type"`
 	StoredAtLT time.Time          `json:"storedAtLT"`
@@ -204,6 +210,10 @@ func (r RealizationRunBase) Validate() error {
 		errs = append(errs, fmt.Errorf("line id must be non-empty"))
 	}
 
+	if r.InvoiceID != nil && *r.InvoiceID == "" {
+		errs = append(errs, fmt.Errorf("invoice id must be non-empty"))
+	}
+
 	if err := r.Type.Validate(); err != nil {
 		errs = append(errs, fmt.Errorf("type: %w", err))
 	}

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

@@ -212,6 +212,7 @@ func (s *CreditThenInvoiceStateMachine) DeleteCharge(ctx context.Context, _ meta
 
 type invoiceCreatedInput struct {
 	LineID          string
+	InvoiceID       string
 	ServicePeriodTo time.Time
 }
 
@@ -220,6 +221,10 @@ func (i invoiceCreatedInput) Validate() error {
 		return fmt.Errorf("line id is required")
 	}
 
+	if i.InvoiceID == "" {
+		return fmt.Errorf("invoice id is required")
+	}
+
 	if i.ServicePeriodTo.IsZero() {
 		return fmt.Errorf("service period to is required")
 	}
@@ -255,6 +260,7 @@ func (s *CreditThenInvoiceStateMachine) startInvoiceCreatedRun(
 		StoredAtLT:         storedAtLT,
 		ServicePeriodTo:    servicePeriodTo,
 		LineID:             lo.ToPtr(input.LineID),
+		InvoiceID:          lo.ToPtr(input.InvoiceID),
 		CreditAllocation:   usagebasedrun.CreditAllocationAvailable,
 		CurrencyCalculator: s.CurrencyCalculator,
 	})

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

@@ -17,7 +17,8 @@ func TestStartInvoiceCreatedRunValidatesInput(t *testing.T) {
 	err := machine.startInvoiceCreatedRun(
 		t.Context(),
 		invoiceCreatedInput{
-			LineID: "line-1",
+			LineID:    "line-1",
+			InvoiceID: "invoice-1",
 		},
 		usagebased.RealizationRunTypePartialInvoice,
 	)

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

@@ -149,6 +149,7 @@ func (e *LineEngine) OnStandardInvoiceCreated(ctx context.Context, input billing
 
 		if err := stateMachine.FireAndActivate(ctx, trigger, invoiceCreatedInput{
 			LineID:          stdLine.ID,
+			InvoiceID:       input.Invoice.ID,
 			ServicePeriodTo: stdLine.Period.To,
 		}); err != nil {
 			return nil, fmt.Errorf("triggering %s for charge[%s]: %w", trigger, stateMachine.GetCharge().ID, err)

openmeter/billing/charges/usagebased/service/run/create.go

@@ -24,6 +24,7 @@ type CreateRatedRunInput struct {
 	StoredAtLT         time.Time
 	ServicePeriodTo    time.Time
 	LineID             *string
+	InvoiceID          *string
 	CreditAllocation   CreditAllocationMode
 	CurrencyCalculator currencyx.Calculator
 	// NoFiatTransactionRequired is set if either there's no fiat-based
@@ -73,6 +74,10 @@ func (i CreateRatedRunInput) Validate() error {
 		return fmt.Errorf("line id if set, must be non-empty")
 	}
 
+	if i.InvoiceID != nil && *i.InvoiceID == "" {
+		return fmt.Errorf("invoice id if set, must be non-empty")
+	}
+
 	if err := i.CreditAllocation.Validate(); err != nil {
 		return fmt.Errorf("credit allocation: %w", err)
 	}
@@ -152,6 +157,7 @@ func (s *Service) CreateRatedRun(ctx context.Context, in CreateRatedRunInput) (C
 		StoredAtLT:                in.StoredAtLT,
 		ServicePeriodTo:           in.ServicePeriodTo,
 		LineID:                    in.LineID,
+		InvoiceID:                 in.InvoiceID,
 		MeteredQuantity:           ratingResult.Quantity,
 		Totals:                    runTotals,
 		NoFiatTransactionRequired: noFiatTransactionRequired,