satellite/{cmd, payments}: add optional command flag to toggle invoice items aggregation

Added new optional '--aggregate' flag for billing.generate-invoices and billing.prepare-invoice-records commands to toggle invoice items aggregation.
Added new explicit command billing.create-aggregated-project-invoice-items which should be used after preparing aggregated invoice records (in cases when invoice generation happens step-by-step).

Change-Id: I04fc0110be5263edb959306d5314a4a1a8eec3ba
This commit is contained in:
Vitalii 2023-11-22 13:33:59 +02:00 committed by Vitalii Shpital
parent d2819522c6
commit 5e3cab29a2
4 changed files with 71 additions and 28 deletions

View File

@ -226,6 +226,9 @@ var (
Long: "Creates stripe invoice line items for stripe customer balances obtained from past invoices and other miscellaneous charges.",
RunE: cmdCreateCustomerBalanceInvoiceItems,
}
aggregate = false
prepareCustomerInvoiceRecordsCmd = &cobra.Command{
Use: "prepare-invoice-records [period]",
Short: "Prepares invoice project records",
@ -240,6 +243,13 @@ var (
Args: cobra.ExactArgs(1),
RunE: cmdCreateCustomerProjectInvoiceItems,
}
createCustomerAggregatedProjectInvoiceItemsCmd = &cobra.Command{
Use: "create-aggregated-project-invoice-items [period]",
Short: "Creates aggregated stripe invoice line items for project charges",
Long: "Creates aggregated stripe invoice line items for not consumed project records.",
Args: cobra.ExactArgs(1),
RunE: cmdCreateAggregatedCustomerProjectInvoiceItems,
}
createCustomerInvoicesCmd = &cobra.Command{
Use: "create-invoices [period]",
Short: "Creates stripe invoices from pending invoice items",
@ -422,9 +432,12 @@ func init() {
billingCmd.AddCommand(setInvoiceStatusCmd)
billingCmd.AddCommand(createCustomerBalanceInvoiceItemsCmd)
billingCmd.AddCommand(prepareCustomerInvoiceRecordsCmd)
prepareCustomerInvoiceRecordsCmd.Flags().BoolVar(&aggregate, "aggregate", false, "Used to enable creation of to be aggregated project records in case users have many projects (more than 83).")
billingCmd.AddCommand(createCustomerProjectInvoiceItemsCmd)
billingCmd.AddCommand(createCustomerAggregatedProjectInvoiceItemsCmd)
billingCmd.AddCommand(createCustomerInvoicesCmd)
billingCmd.AddCommand(generateCustomerInvoicesCmd)
generateCustomerInvoicesCmd.Flags().BoolVar(&aggregate, "aggregate", false, "Used to enable invoice items aggregation in case users have many projects (more than 83).")
billingCmd.AddCommand(finalizeCustomerInvoicesCmd)
billingCmd.AddCommand(payInvoicesWithTokenCmd)
billingCmd.AddCommand(payAllInvoicesCmd)
@ -843,7 +856,7 @@ func cmdPrepareCustomerInvoiceRecords(cmd *cobra.Command, args []string) (err er
}
return runBillingCmd(ctx, func(ctx context.Context, payments *stripe.Service, _ satellite.DB) error {
return payments.PrepareInvoiceProjectRecords(ctx, periodStart)
return payments.PrepareInvoiceProjectRecords(ctx, periodStart, aggregate)
})
}
@ -860,6 +873,19 @@ func cmdCreateCustomerProjectInvoiceItems(cmd *cobra.Command, args []string) (er
})
}
func cmdCreateAggregatedCustomerProjectInvoiceItems(cmd *cobra.Command, args []string) (err error) {
ctx, _ := process.Ctx(cmd)
periodStart, err := parseYearMonth(args[0])
if err != nil {
return err
}
return runBillingCmd(ctx, func(ctx context.Context, payments *stripe.Service, _ satellite.DB) error {
return payments.InvoiceApplyToBeAggregatedProjectRecords(ctx, periodStart)
})
}
func cmdCreateCustomerInvoices(cmd *cobra.Command, args []string) (err error) {
ctx, _ := process.Ctx(cmd)
@ -882,7 +908,7 @@ func cmdGenerateCustomerInvoices(cmd *cobra.Command, args []string) (err error)
}
return runBillingCmd(ctx, func(ctx context.Context, payments *stripe.Service, _ satellite.DB) error {
return payments.GenerateInvoices(ctx, periodStart)
return payments.GenerateInvoices(ctx, periodStart, aggregate)
})
}

View File

@ -709,7 +709,7 @@ func TestProjectCheckUsage_lastMonthUnappliedInvoice(t *testing.T) {
planet.Satellites[0].API.Payments.StripeService.SetNow(func() time.Time {
return oneMonthAhead
})
err = planet.Satellites[0].API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, now)
err = planet.Satellites[0].API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, now, false)
require.NoError(t, err)
req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://"+address.String()+"/api/projects/%s/usage", projectID), nil)

View File

@ -135,7 +135,7 @@ func (service *Service) Accounts() payments.Accounts {
}
// PrepareInvoiceProjectRecords iterates through all projects and creates invoice records if none exist.
func (service *Service) PrepareInvoiceProjectRecords(ctx context.Context, period time.Time) (err error) {
func (service *Service) PrepareInvoiceProjectRecords(ctx context.Context, period time.Time, shouldAggregate bool) (err error) {
defer mon.Task()(&ctx)(&err)
now := service.nowFn().UTC()
@ -164,7 +164,7 @@ func (service *Service) PrepareInvoiceProjectRecords(ctx context.Context, period
}
numberOfCustomers += len(customersPage.Customers)
records, err := service.processCustomers(ctx, customersPage.Customers, start, end)
records, err := service.processCustomers(ctx, customersPage.Customers, shouldAggregate, start, end)
if err != nil {
return Error.Wrap(err)
}
@ -175,7 +175,7 @@ func (service *Service) PrepareInvoiceProjectRecords(ctx context.Context, period
return nil
}
func (service *Service) processCustomers(ctx context.Context, customers []Customer, start, end time.Time) (int, error) {
func (service *Service) processCustomers(ctx context.Context, customers []Customer, shouldAggregate bool, start, end time.Time) (int, error) {
var regularRecords []CreateProjectRecord
var recordsToAggregate []CreateProjectRecord
for _, customer := range customers {
@ -197,7 +197,7 @@ func (service *Service) processCustomers(ctx context.Context, customers []Custom
// We generate 3 invoice items for each user project which means,
// we can support only 83 projects in a single invoice (249 invoice items).
if len(projects) > 83 {
if shouldAggregate && len(projects) > 83 {
recordsToAggregate = append(recordsToAggregate, records...)
} else {
regularRecords = append(regularRecords, records...)
@ -209,12 +209,18 @@ func (service *Service) processCustomers(ctx context.Context, customers []Custom
return 0, err
}
err = service.db.ProjectRecords().CreateToBeAggregated(ctx, recordsToAggregate, start, end)
if err != nil {
return 0, err
recordsCount := len(regularRecords)
if shouldAggregate {
err = service.db.ProjectRecords().CreateToBeAggregated(ctx, recordsToAggregate, start, end)
if err != nil {
return 0, err
}
recordsCount += len(recordsToAggregate)
}
return len(recordsToAggregate) + len(regularRecords), nil
return recordsCount, nil
}
// createProjectRecords creates invoice project record if none exists.
@ -1253,24 +1259,35 @@ func (service *Service) CreateBalanceInvoiceItems(ctx context.Context) (err erro
// GenerateInvoices performs tasks necessary to generate Stripe invoices.
// This is equivalent to invoking PrepareInvoiceProjectRecords, InvoiceApplyProjectRecords,
// and CreateInvoices in order.
func (service *Service) GenerateInvoices(ctx context.Context, period time.Time) (err error) {
func (service *Service) GenerateInvoices(ctx context.Context, period time.Time, shouldAggregate bool) (err error) {
defer mon.Task()(&ctx)(&err)
for _, subFn := range []struct {
Description string
Exec func(context.Context, time.Time) error
}{
{"Preparing invoice project records", service.PrepareInvoiceProjectRecords},
{"Applying invoice project records", service.InvoiceApplyProjectRecords},
{"Applying to be aggregated invoice project records", service.InvoiceApplyToBeAggregatedProjectRecords},
{"Creating invoices", service.CreateInvoices},
} {
service.log.Info(subFn.Description)
if err := subFn.Exec(ctx, period); err != nil {
service.log.Info("Preparing invoice project records")
err = service.PrepareInvoiceProjectRecords(ctx, period, shouldAggregate)
if err != nil {
return err
}
service.log.Info("Applying invoice project records")
err = service.InvoiceApplyProjectRecords(ctx, period)
if err != nil {
return err
}
if shouldAggregate {
service.log.Info("Applying to be aggregated invoice project records")
err = service.InvoiceApplyToBeAggregatedProjectRecords(ctx, period)
if err != nil {
return err
}
}
service.log.Info("Creating invoices")
err = service.CreateInvoices(ctx, period)
if err != nil {
return err
}
return nil
}

View File

@ -450,7 +450,7 @@ func TestService_InvoiceElementsProcessing(t *testing.T) {
satellite.API.Payments.StripeService.SetNow(func() time.Time {
return time.Date(period.Year(), period.Month()+1, 1, 0, 0, 0, 0, time.UTC)
})
err := satellite.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period)
err := satellite.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period, false)
require.NoError(t, err)
start := time.Date(period.Year(), period.Month(), 1, 0, 0, 0, 0, time.UTC)
@ -532,7 +532,7 @@ func TestService_InvoiceUserWithManyProjects(t *testing.T) {
require.Nil(t, projectRecord)
}
err = payments.StripeService.PrepareInvoiceProjectRecords(ctx, period)
err = payments.StripeService.PrepareInvoiceProjectRecords(ctx, period, false)
require.NoError(t, err)
for i := 0; i < len(projects); i++ {
@ -703,7 +703,7 @@ func TestService_ProjectsWithMembers(t *testing.T) {
satellite.API.Payments.StripeService.SetNow(func() time.Time {
return time.Date(period.Year(), period.Month()+1, 1, 0, 0, 0, 0, time.UTC)
})
err := satellite.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period)
err := satellite.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period, false)
require.NoError(t, err)
start := time.Date(period.Year(), period.Month(), 1, 0, 0, 0, 0, time.UTC)
@ -1318,7 +1318,7 @@ func TestService_GenerateInvoice(t *testing.T) {
99)
}
require.NoError(t, payments.StripeService.GenerateInvoices(ctx, start))
require.NoError(t, payments.StripeService.GenerateInvoices(ctx, start, false))
// ensure project record was generated
err = satellite.DB.StripeCoinPayments().ProjectRecords().Check(ctx, proj.ID, start, end)
@ -1477,7 +1477,7 @@ func TestProjectUsagePrice(t *testing.T) {
pb.PieceAction_GET, memory.TB.Int64(), 0, period)
require.NoError(t, err)
err = sat.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period)
err = sat.API.Payments.StripeService.PrepareInvoiceProjectRecords(ctx, period, false)
require.NoError(t, err)
err = sat.API.Payments.StripeService.InvoiceApplyProjectRecords(ctx, period)