[usage] implement CancelSubscription

Author: svenefftinge

components/usage/pkg/apiv1/billing.go

@@ -22,16 +22,18 @@ import (
 	"gorm.io/gorm"
 )
 
-func NewBillingService(stripeClient *stripe.Client, conn *gorm.DB) *BillingService {
+func NewBillingService(stripeClient *stripe.Client, conn *gorm.DB, ccManager *db.CostCenterManager) *BillingService {
 	return &BillingService{
 		stripeClient: stripeClient,
 		conn:         conn,
+		ccManager:    ccManager,
 	}
 }
 
 type BillingService struct {
 	conn         *gorm.DB
 	stripeClient *stripe.Client
+	ccManager    *db.CostCenterManager
 
 	v1.UnimplementedBillingServiceServer
 }
@@ -76,19 +78,18 @@ func (s *BillingService) FinalizeInvoice(ctx context.Context, in *v1.FinalizeInv
 		return nil, status.Errorf(codes.Internal, "Failed to retrieve subscription details from invoice.")
 	}
 
-	teamID, found := subscription.Metadata[stripe.AttributionIDMetadataKey]
+	attrID, found := subscription.Metadata[stripe.AttributionIDMetadataKey]
 	if !found {
 		logger.Error("Failed to find teamID from subscription metadata.")
 		return nil, status.Errorf(codes.Internal, "Failed to extra teamID from Stripe subscription.")
 	}
-	logger = logger.WithField("team_id", teamID)
+	logger = logger.WithField("attribution_id", attrID)
 
 	// To support individual `user`s, we'll need to also extract the `userId` from metadata here and handle separately.
-	attributionID := db.NewTeamAttributionID(teamID)
+	attributionID := db.NewTeamAttributionID(attrID)
 	finalizedAt := time.Unix(invoice.StatusTransitions.FinalizedAt, 0)
 
 	logger = logger.
-		WithField("attribution_id", attributionID).
 		WithField("invoice_finalized_at", finalizedAt)
 
 	if invoice.Lines == nil || len(invoice.Lines.Data) == 0 {
@@ -126,6 +127,29 @@ func (s *BillingService) FinalizeInvoice(ctx context.Context, in *v1.FinalizeInv
 	return &v1.FinalizeInvoiceResponse{}, nil
 }
 
+func (s *BillingService) CancelSubscription(ctx context.Context, in *v1.CancelSubscriptionRequest) (*v1.CancelSubscriptionResponse, error) {
+	logger := log.WithField("subscription_id", in.GetSubscriptionId())
+	if in.GetSubscriptionId() == "" {
+		return nil, status.Errorf(codes.InvalidArgument, "subscriptionId is required")
+	}
+
+	attributionID, err := s.stripeClient.GetAttributionIdForSubscriptionId(ctx, in.GetSubscriptionId())
+	if err != nil {
+		return nil, err
+	}
+	logger.Infof("Subscription ended. Setting cost center back to free.")
+	costCenter, err := s.ccManager.GetOrCreateCostCenter(ctx, attributionID)
+	if err != nil {
+		return nil, err
+	}
+	costCenter.BillingStrategy = db.CostCenter_Other
+	_, err = s.ccManager.UpdateCostCenter(ctx, costCenter)
+	if err != nil {
+		return nil, err
+	}
+	return &v1.CancelSubscriptionResponse{}, nil
+}
+
 func (s *BillingService) GetUpcomingInvoice(ctx context.Context, in *v1.GetUpcomingInvoiceRequest) (*v1.GetUpcomingInvoiceResponse, error) {
 	if in.GetTeamId() == "" && in.GetUserId() == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "teamId or userId is required")

components/usage/pkg/db/cost_center.go

@@ -112,13 +112,16 @@ func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter Cos
 
 	now := time.Now()
 
+	// we always update the creationTime
+	costCenter.CreationTime = NewVarcharTime(now)
 	// we don't allow setting the creationTime or the nextBillingTime from outside
 	costCenter.CreationTime = existingCostCenter.CreationTime
 	costCenter.NextBillingTime = existingCostCenter.NextBillingTime
 
 	// Do we have a billing strategy update?
 	if costCenter.BillingStrategy != existingCostCenter.BillingStrategy {
-		if existingCostCenter.BillingStrategy == CostCenter_Other {
+		switch costCenter.BillingStrategy {
+		case CostCenter_Stripe:
 			// moving to stripe -> let's run a finalization
 			finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, costCenter.ID)
 			if err != nil {
@@ -130,12 +133,21 @@ func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter Cos
 					return CostCenter{}, err
 				}
 			}
+			// we don't manage stripe billing cycle
+			costCenter.NextBillingTime = VarcharTime{}
+
+		case CostCenter_Other:
+			// cancelled from stripe reset the spending limit
+			if costCenter.ID.IsEntity(AttributionEntity_Team) {
+				costCenter.SpendingLimit = c.cfg.ForTeams
+			} else {
+				costCenter.SpendingLimit = c.cfg.ForUsers
+			}
+			// see you next month
+			costCenter.NextBillingTime = NewVarcharTime(now.AddDate(0, 1, 0))
 		}
-		c.updateNextBillingTime(&costCenter, now)
 	}
 
-	// we update the creationTime
-	costCenter.CreationTime = NewVarcharTime(now)
 	db := c.conn.Save(&costCenter)
 	if db.Error != nil {
 		return CostCenter{}, fmt.Errorf("failed to save cost center for attributionID %s: %w", costCenter.ID, db.Error)
@@ -163,8 +175,3 @@ func (c *CostCenterManager) ComputeInvoiceUsageRecord(ctx context.Context, attri
 		Draft:         false,
 	}, nil
 }
-
-func (c *CostCenterManager) updateNextBillingTime(costCenter *CostCenter, now time.Time) {
-	nextMonth := NewVarcharTime(time.Now().AddDate(0, 1, 0))
-	costCenter.NextBillingTime = nextMonth
-}

components/usage/pkg/db/cost_center_test.go

@@ -83,20 +83,28 @@ func TestCostCenterManager_UpdateCostCenter(t *testing.T) {
 func TestSaveCostCenterMovedToStripe(t *testing.T) {
 	conn := dbtest.ConnectForTests(t)
 	mnr := db.NewCostCenterManager(conn, db.DefaultSpendingLimit{
-		ForTeams: 0,
+		ForTeams: 20,
 		ForUsers: 500,
 	})
 	team := db.NewTeamAttributionID(uuid.New().String())
 	cleanUp(t, conn, team)
 	teamCC, err := mnr.GetOrCreateCostCenter(context.Background(), team)
 	require.NoError(t, err)
-	require.Equal(t, int32(0), teamCC.SpendingLimit)
+	require.Equal(t, int32(20), teamCC.SpendingLimit)
 
 	teamCC.BillingStrategy = db.CostCenter_Stripe
-	newTeamCC, err := mnr.UpdateCostCenter(context.Background(), teamCC)
+	teamCC.SpendingLimit = 400050
+	teamCC, err = mnr.UpdateCostCenter(context.Background(), teamCC)
+	require.NoError(t, err)
+	require.Equal(t, db.CostCenter_Stripe, teamCC.BillingStrategy)
+	require.Equal(t, db.VarcharTime{}, teamCC.NextBillingTime)
+	require.Equal(t, int32(400050), teamCC.SpendingLimit)
+
+	teamCC.BillingStrategy = db.CostCenter_Other
+	teamCC, err = mnr.UpdateCostCenter(context.Background(), teamCC)
 	require.NoError(t, err)
-	require.Equal(t, db.CostCenter_Stripe, newTeamCC.BillingStrategy)
-	require.Equal(t, newTeamCC.CreationTime.Time().AddDate(0, 1, 0).Truncate(time.Second), newTeamCC.NextBillingTime.Time().Truncate(time.Second))
+	require.Equal(t, teamCC.CreationTime.Time().AddDate(0, 1, 0).Truncate(time.Second), teamCC.NextBillingTime.Time().Truncate(time.Second))
+	require.Equal(t, int32(20), teamCC.SpendingLimit)
 }
 
 func cleanUp(t *testing.T, conn *gorm.DB, attributionIds ...db.AttributionID) {

components/usage/pkg/server/server.go

@@ -156,7 +156,7 @@ func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *s
 	if stripeClient == nil {
 		v1.RegisterBillingServiceServer(srv.GRPC(), &apiv1.BillingServiceNoop{})
 	} else {
-		v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient, conn))
+		v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient, conn, ccManager))
 	}
 	return nil
 }

components/usage/pkg/stripe/stripe.go

@@ -8,10 +8,11 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/gitpod-io/gitpod/usage/pkg/db"
 	"os"
 	"strings"
 
+	"github.com/gitpod-io/gitpod/usage/pkg/db"
+
 	"github.com/gitpod-io/gitpod/common-go/log"
 	"github.com/stripe/stripe-go/v72"
 	"github.com/stripe/stripe-go/v72/client"
@@ -244,6 +245,19 @@ func (c *Client) GetInvoice(ctx context.Context, invoiceID string) (*stripe.Invo
 	return invoice, nil
 }
 
+func (c *Client) GetAttributionIdForSubscriptionId(ctx context.Context, subscriptionID string) (db.AttributionID, error) {
+	subscription, err := c.sc.Subscriptions.Get(subscriptionID, nil)
+	if err != nil {
+		return "", fmt.Errorf("failed to search for subscription (%s): %w", subscriptionID, err)
+	}
+	attributionIDRaw := subscription.Customer.Metadata[AttributionIDMetadataKey]
+	attributionID, err := db.ParseAttributionID(attributionIDRaw)
+	if err != nil {
+		return "", fmt.Errorf("failed to fetch AttributionID for subscription (%s): %w", subscriptionID, err)
+	}
+	return attributionID, nil
+}
+
 // queriesForCustomersWithAttributionIDs constructs Stripe query strings to find the Stripe Customer for each teamId
 // It returns multiple queries, each being a big disjunction of subclauses so that we can process multiple teamIds in one query.
 // `clausesPerQuery` is a limit enforced by the Stripe API.