Skip to content

[usage] Use stripe clients rather than a singleton in the usage controller #10854

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 27, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 19 additions & 6 deletions components/usage/cmd/run.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,17 @@
package cmd

import (
"encoding/json"
"net"
"os"
"time"

"github.com/gitpod-io/gitpod/common-go/baseserver"
"github.com/gitpod-io/gitpod/common-go/log"
"github.com/gitpod-io/gitpod/usage/pkg/controller"
"github.com/gitpod-io/gitpod/usage/pkg/db"
"github.com/gitpod-io/gitpod/usage/pkg/stripe"
"github.com/spf13/cobra"
"net"
"os"
"time"
)

func init() {
Expand Down Expand Up @@ -47,11 +49,22 @@ func run() *cobra.Command {
var billingController controller.BillingController = &controller.NoOpBillingController{}

if apiKeyFile != "" {
err = stripe.Authenticate(apiKeyFile)
bytes, err := os.ReadFile(apiKeyFile)
if err != nil {
log.WithError(err).Fatal("Failed to read Stripe API keys.")
}

var config stripe.ClientConfig
err = json.Unmarshal(bytes, &config)
if err != nil {
log.WithError(err).Fatal("Failed to unmarshal Stripe API keys.")
}

c, err := stripe.New(config)
if err != nil {
log.WithError(err).Fatal("Failed to initialize stripe client.")
log.WithError(err).Fatal("Failed to initialize Stripe client.")
}
billingController = &controller.StripeBillingController{}
billingController = controller.NewStripeBillingController(c)
}

ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, billingController))
Expand Down
11 changes: 9 additions & 2 deletions components/usage/pkg/controller/billing.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,16 +11,23 @@ type BillingController interface {
}

type NoOpBillingController struct{}
type StripeBillingController struct{}

func (b *NoOpBillingController) Reconcile(report []TeamUsage) {}

type StripeBillingController struct {
sc *stripe.Client
}

func NewStripeBillingController(sc *stripe.Client) *StripeBillingController {
return &StripeBillingController{sc: sc}
}

func (b *StripeBillingController) Reconcile(report []TeamUsage) {
// Convert the usage report to sum all entries for the same team.
var summedReport = make(map[string]int64)
for _, usageEntry := range report {
summedReport[usageEntry.TeamID] += usageEntry.WorkspaceSeconds
}

stripe.UpdateUsage(summedReport)
b.sc.UpdateUsage(summedReport)
}
35 changes: 12 additions & 23 deletions components/usage/pkg/stripe/stripe.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,43 +5,32 @@
package stripe

import (
"encoding/json"
"fmt"
"math"
"os"
"strings"

"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"
"github.com/stripe/stripe-go/v72/client"
)

type stripeKeys struct {
type Client struct {
sc *client.API
}

type ClientConfig struct {
PublishableKey string `json:"publishableKey"`
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 {
return err
}

var stripeKeys stripeKeys
err = json.Unmarshal(bytes, &stripeKeys)
if err != nil {
return err
}

stripe.Key = stripeKeys.SecretKey
return nil
// New authenticates a Stripe client using the provided config
func New(config ClientConfig) (*Client, error) {
return &Client{sc: client.New(config.SecretKey, nil)}, nil
}

// 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 {
func (c *Client) UpdateUsage(usageForTeam map[string]int64) error {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For follow-up, I'd recommend returnig a struct which contains a summary of the succesful and failed updates. That way, you can also better test that the right number of records got updated. For an example, see the Usage Reconciler

teamIds := make([]string, 0, len(usageForTeam))
for k := range usageForTeam {
teamIds = append(teamIds, k)
Expand All @@ -56,7 +45,7 @@ func UpdateUsage(usageForTeam map[string]int64) error {
Expand: []*string{stripe.String("data.subscriptions")},
},
}
iter := customer.Search(params)
iter := c.sc.Customers.Search(params)
for iter.Next() {
customer := iter.Customer()
log.Infof("found customer %q for teamId %q", customer.Name, customer.Metadata["teamId"])
Expand All @@ -77,7 +66,7 @@ func UpdateUsage(usageForTeam map[string]int64) error {

subscriptionItemId := subscription.Items.Data[0].ID
log.Infof("registering usage against subscriptionItem %q", subscriptionItemId)
_, err := usagerecord.New(&stripe.UsageRecordParams{
_, err := c.sc.UsageRecords.New(&stripe.UsageRecordParams{
SubscriptionItem: stripe.String(subscriptionItemId),
Quantity: stripe.Int64(creditsUsed),
})
Expand Down