[usage] update cost center on stripe cancellation

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 {
@@ -123,6 +124,19 @@ func (s *BillingService) FinalizeInvoice(ctx context.Context, in *v1.FinalizeInv
 
 	logger.WithField("usage_id", usage.ID).Infof("Inserted usage record into database for %d credits against %s attribution", creditsOnInvoice, attributionID)
 
+	if subscription.EndedAt < time.Now().Unix() {
+		logger.Infof("Subscription ended. Setting back to default.")
+		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.FinalizeInvoiceResponse{}, nil
 }
 

components/usage/pkg/db/cost_center.go

@@ -112,13 +112,15 @@ func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter Cos
 
 	now := time.Now()
 
+	// we 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 {
+		if costCenter.BillingStrategy == CostCenter_Stripe {
 			// moving to stripe -> let's run a finalization
 			finalizationUsage, err := c.ComputeInvoiceUsageRecord(ctx, costCenter.ID)
 			if err != nil {
@@ -130,12 +132,20 @@ func (c *CostCenterManager) UpdateCostCenter(ctx context.Context, costCenter Cos
 					return CostCenter{}, err
 				}
 			}
+			// we don't manage stripe billing cycle
+			costCenter.NextBillingTime = VarcharTime{}
+		} else if costCenter.BillingStrategy == 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 +173,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
 }