diff --git a/components/usage/pkg/controller/reconciler.go b/components/usage/pkg/controller/reconciler.go index 3c96bd29abbdb5..7b9b53f5dbf055 100644 --- a/components/usage/pkg/controller/reconciler.go +++ b/components/usage/pkg/controller/reconciler.go @@ -132,18 +132,13 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time. } func submitUsageReport(report []TeamUsage) { - var teamIdSet = make(map[string]bool) - - // Convert the usage report into a set of teamIds occurring in the report + // Convert the usage report to sum all entries for the same team. + var summedReport = make(map[string]int64) for _, usageEntry := range report { - teamIdSet[usageEntry.TeamID] = true - } - teamIds := make([]string, 0, len(teamIdSet)) - for k := range teamIdSet { - teamIds = append(teamIds, k) + summedReport[usageEntry.TeamID] += usageEntry.WorkspaceSeconds } - stripe.FindCustomersForTeamIds(teamIds) + stripe.UpdateUsage(summedReport) } func generateUsageReport(teams []teamWithWorkspaces, maxStopTime time.Time) ([]TeamUsage, error) { diff --git a/components/usage/pkg/stripe/stripe.go b/components/usage/pkg/stripe/stripe.go index c9e88e766e1e18..ee16665c31d72b 100644 --- a/components/usage/pkg/stripe/stripe.go +++ b/components/usage/pkg/stripe/stripe.go @@ -13,6 +13,7 @@ import ( "github.com/gitpod-io/gitpod/common-go/log" "github.com/stripe/stripe-go/v72" "github.com/stripe/stripe-go/v72/customer" + "github.com/stripe/stripe-go/v72/usagerecord" ) type stripeKeys struct { @@ -20,6 +21,7 @@ type stripeKeys struct { SecretKey string `json:"secretKey"` } +// Authenticate authenticates the Stripe client using a provided file containing a Stripe secret key. func Authenticate(apiKeyFile string) error { bytes, err := os.ReadFile(apiKeyFile) if err != nil { @@ -36,19 +38,54 @@ func Authenticate(apiKeyFile string) error { return nil } -// FindCustomersForTeamIds queries the stripe API to find all customers with a teamId in `teamIds`. -func FindCustomersForTeamIds(teamIds []string) { +// UpdateUsage updates teams' Stripe subscriptions with usage data +// `usageForTeam` is a map from team name to total workspace seconds used within a billing period. +func UpdateUsage(usageForTeam map[string]int64) error { + teamIds := make([]string, 0, len(usageForTeam)) + for k := range usageForTeam { + teamIds = append(teamIds, k) + } queries := queriesForCustomersWithTeamIds(teamIds) for _, query := range queries { log.Infof("about to make query %q", query) - params := &stripe.CustomerSearchParams{SearchParams: stripe.SearchParams{Query: query}} + params := &stripe.CustomerSearchParams{ + SearchParams: stripe.SearchParams{ + Query: query, + Expand: []*string{stripe.String("data.subscriptions")}, + }, + } iter := customer.Search(params) for iter.Next() { customer := iter.Customer() log.Infof("found customer %q for teamId %q", customer.Name, customer.Metadata["teamId"]) + subscriptions := customer.Subscriptions.Data + if len(subscriptions) != 1 { + log.Errorf("customer has an unexpected number of subscriptions (expected 1, got %d)", len(subscriptions)) + continue + } + subscription := customer.Subscriptions.Data[0] + + log.Infof("customer has subscription: %q", subscription.ID) + if len(subscription.Items.Data) != 1 { + log.Errorf("this subscription has an unexpected number of subscriptionItems (expected 1, got %d)", len(subscription.Items.Data)) + continue + } + + creditsUsed := workspaceSecondsToCredits(usageForTeam[customer.Metadata["teamId"]]) + + subscriptionItemId := subscription.Items.Data[0].ID + log.Infof("registering usage against subscriptionItem %q", subscriptionItemId) + _, err := usagerecord.New(&stripe.UsageRecordParams{ + SubscriptionItem: stripe.String(subscriptionItemId), + Quantity: stripe.Int64(creditsUsed), + }) + if err != nil { + log.WithError(err).Errorf("failed to register usage for customer %q", customer.Name) + } } } + return nil } // queriesForCustomersWithTeamIds constructs Stripe query strings to find the Stripe Customer for each teamId @@ -72,3 +109,8 @@ func queriesForCustomersWithTeamIds(teamIds []string) []string { return queries } + +// workspaceSecondsToCredits converts seconds (of workspace usage) into Stripe credits. +func workspaceSecondsToCredits(seconds int64) int64 { + return (seconds + 59) / 60 +} diff --git a/components/usage/pkg/stripe/stripe_test.go b/components/usage/pkg/stripe/stripe_test.go index c628e372577a59..a0a45cfa267454 100644 --- a/components/usage/pkg/stripe/stripe_test.go +++ b/components/usage/pkg/stripe/stripe_test.go @@ -83,3 +83,50 @@ func TestCustomerQueriesForTeamIds_MultipleQueries(t *testing.T) { }) } } + +func TestWorkspaceSecondsToCreditsCalcuation(t *testing.T) { + testCases := []struct { + Name string + Seconds int64 + ExpectedCredits int64 + }{ + { + Name: "0 seconds", + Seconds: 0, + ExpectedCredits: 0, + }, + { + Name: "1 second", + Seconds: 1, + ExpectedCredits: 1, + }, + { + Name: "60 seconds", + Seconds: 60, + ExpectedCredits: 1, + }, + { + Name: "61 seconds", + Seconds: 61, + ExpectedCredits: 2, + }, + { + Name: "90 seconds", + Seconds: 90, + ExpectedCredits: 2, + }, + { + Name: "1 hour", + Seconds: 3600, + ExpectedCredits: 60, + }, + } + + for _, tc := range testCases { + t.Run(tc.Name, func(t *testing.T) { + actualCredits := workspaceSecondsToCredits(tc.Seconds) + + require.Equal(t, tc.ExpectedCredits, actualCredits) + }) + } +}