Clean up billing customer creation and improve error handling (#4559)

api/server/handlers/billing/list.go

@@ -1,6 +1,7 @@
 package billing
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -65,87 +66,155 @@ func (c *CheckPaymentEnabledHandler) ServeHTTP(w http.ResponseWriter, r *http.Re
 	defer span.End()
 
 	proj, _ := ctx.Value(types.ProjectScope).(*models.Project)
+	currentUser, _ := ctx.Value(types.UserScope).(*models.User)
 
-	// Get project roles
-	roles, err := c.Repo().Project().ListProjectRolesOrdered(proj.ID)
+	err := c.ensureBillingSetup(ctx, proj, currentUser)
 	if err != nil {
-		err = telemetry.Error(ctx, span, err, "error listing project roles")
+		err := telemetry.Error(ctx, span, err, "error ensuring billing setup")
 		c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
 		return
 	}
 
-	// Get the project admin user
-	var adminUser *models.User
-	for _, role := range roles {
-		if role.Kind == types.RoleAdmin {
-			adminUser, err = c.Repo().User().ReadUser(role.UserID)
-			if err != nil {
-				err = telemetry.Error(ctx, span, err, "error reading user")
-				c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
-				return
-			}
-			break
-		}
+	telemetry.WithAttributes(span,
+		telemetry.AttributeKV{Key: "project-id", Value: proj.ID},
+		telemetry.AttributeKV{Key: "customer-id", Value: proj.BillingID},
+	)
+
+	paymentEnabled, err := c.Config().BillingManager.StripeClient.CheckPaymentEnabled(ctx, proj.BillingID)
+	if err != nil {
+		err := telemetry.Error(ctx, span, err, "error checking if payment enabled")
+		c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
+		return
 	}
 
 	telemetry.WithAttributes(span,
-		telemetry.AttributeKV{Key: "admin-user-id", Value: adminUser.ID},
-		telemetry.AttributeKV{Key: "admin-user-email", Value: adminUser.Email},
+		telemetry.AttributeKV{Key: "payment-enabled", Value: paymentEnabled},
 	)
 
-	// Create billing customer for project and set the billing ID if it doesn't exist
-	var shouldUpdate bool
-	if proj.BillingID == "" {
-		billingID, err := c.Config().BillingManager.StripeClient.CreateCustomer(ctx, adminUser.Email, proj.ID, proj.Name)
+	c.WriteResult(w, r, paymentEnabled)
+}
+
+func (c *CheckPaymentEnabledHandler) ensureBillingSetup(ctx context.Context, proj *models.Project, user *models.User) (err error) {
+	ctx, span := telemetry.NewSpan(ctx, "ensure-billing-setup")
+	defer span.End()
+
+	telemetry.WithAttributes(span,
+		telemetry.AttributeKV{Key: "billing-id", Value: proj.BillingID},
+		telemetry.AttributeKV{Key: "usage-id", Value: proj.UsageID},
+	)
+
+	if proj.BillingID == "" || proj.UsageID == uuid.Nil {
+		adminUser, err := c.getAdminUser(ctx, proj.ID)
 		if err != nil {
-			err = telemetry.Error(ctx, span, err, "error creating billing customer")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
-			return
+			return telemetry.Error(ctx, span, err, "error getting admin user")
 		}
-		proj.BillingID = billingID
-		shouldUpdate = true
 
-		telemetry.WithAttributes(span,
-			telemetry.AttributeKV{Key: "billing-id", Value: proj.BillingID},
-		)
-	}
+		// If the admin user is not found, use the current user as last resort
+		if adminUser == nil {
+			adminUser = user
+		}
 
-	if c.Config().BillingManager.MetronomeEnabled && proj.GetFeatureFlag(models.MetronomeEnabled, c.Config().LaunchDarklyClient) && proj.UsageID == uuid.Nil {
-		customerID, customerPlanID, err := c.Config().BillingManager.MetronomeClient.CreateCustomerWithPlan(ctx, adminUser.Email, proj.Name, proj.ID, proj.BillingID, proj.EnableSandbox)
+		// Create billing customer for project and set the billing ID if it doesn't exist
+		shouldUpdateBilling, err := c.ensureStripeCustomerExists(ctx, adminUser.Email, proj)
 		if err != nil {
-			err = telemetry.Error(ctx, span, err, "error creating Metronome customer")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
+			return telemetry.Error(ctx, span, err, "error ensuring Stripe customer exists")
+		}
+
+		// Create usage customer for project and set the usage ID if it doesn't exist
+		shouldUpdateUsage, err := c.ensureMetronomeCustomerExists(ctx, adminUser.Email, proj)
+		if err != nil {
+			return telemetry.Error(ctx, span, err, "error ensuring Metronome customer exists")
+		}
+
+		if !shouldUpdateBilling && !shouldUpdateUsage {
+			return nil
+		}
+
+		_, err = c.Repo().Project().UpdateProject(proj)
+		if err != nil {
+			return telemetry.Error(ctx, span, err, "error updating project")
 		}
-		proj.UsageID = customerID
-		proj.UsagePlanID = customerPlanID
-		shouldUpdate = true
-
-		telemetry.WithAttributes(span,
-			telemetry.AttributeKV{Key: "usage-id", Value: proj.UsageID},
-			telemetry.AttributeKV{Key: "usage-plan-id", Value: proj.UsagePlanID},
-		)
 	}
 
-	if shouldUpdate {
-		_, err := c.Repo().Project().UpdateProject(proj)
+	return nil
+}
+
+func (c *CheckPaymentEnabledHandler) getAdminUser(ctx context.Context, projectID uint) (adminUser *models.User, err error) {
+	ctx, span := telemetry.NewSpan(ctx, "get-project-admin-role")
+	defer span.End()
+
+	// Get project roles
+	roles, err := c.Repo().Project().ListProjectRolesOrdered(projectID)
+	if err != nil {
+		return adminUser, telemetry.Error(ctx, span, err, "error listing project roles")
+	}
+
+	// Get the project admin user
+	for _, role := range roles {
+		if role.Kind != types.RoleAdmin {
+			continue
+		}
+
+		adminUser, err = c.Repo().User().ReadUser(role.UserID)
 		if err != nil {
-			err := telemetry.Error(ctx, span, err, "error updating project")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
-			return
+			// If the user is not found, continue to the next role
+			continue
 		}
+		break
 	}
 
-	paymentEnabled, err := c.Config().BillingManager.StripeClient.CheckPaymentEnabled(ctx, proj.BillingID)
+	telemetry.WithAttributes(span,
+		telemetry.AttributeKV{Key: "admin-user-id", Value: adminUser.ID},
+		telemetry.AttributeKV{Key: "admin-user-email", Value: adminUser.Email},
+	)
+
+	return adminUser, nil
+}
+
+func (c *CheckPaymentEnabledHandler) ensureStripeCustomerExists(ctx context.Context, adminUserEmail string, proj *models.Project) (shouldUpdate bool, err error) {
+	ctx, span := telemetry.NewSpan(ctx, "ensure-stripe-customer-exists")
+	defer span.End()
+
+	if proj.BillingID != "" {
+		return false, nil
+	}
+
+	billingID, err := c.Config().BillingManager.StripeClient.CreateCustomer(ctx, adminUserEmail, proj.ID, proj.Name)
 	if err != nil {
-		err := telemetry.Error(ctx, span, err, "error checking if payment enabled")
-		c.HandleAPIError(w, r, apierrors.NewErrInternal(fmt.Errorf("error checking if payment enabled: %w", err)))
-		return
+		return false, telemetry.Error(ctx, span, err, "error creating billing customer")
 	}
 
 	telemetry.WithAttributes(span,
-		telemetry.AttributeKV{Key: "project-id", Value: proj.ID},
-		telemetry.AttributeKV{Key: "customer-id", Value: proj.BillingID},
+		telemetry.AttributeKV{Key: "billing-id", Value: proj.BillingID},
 	)
 
-	c.WriteResult(w, r, paymentEnabled)
+	proj.BillingID = billingID
+	return true, nil
+}
+
+func (c *CheckPaymentEnabledHandler) ensureMetronomeCustomerExists(ctx context.Context, adminUserEmail string, proj *models.Project) (shouldUpdate bool, err error) {
+	ctx, span := telemetry.NewSpan(ctx, "ensure-metronome-customer-exists")
+	defer span.End()
+
+	if proj.UsageID != uuid.Nil {
+		return false, nil
+	}
+
+	if !c.Config().BillingManager.MetronomeEnabled || !proj.GetFeatureFlag(models.MetronomeEnabled, c.Config().LaunchDarklyClient) || proj.UsageID != uuid.Nil {
+		return false, nil
+	}
+
+	customerID, customerPlanID, err := c.Config().BillingManager.MetronomeClient.CreateCustomerWithPlan(ctx, adminUserEmail, proj.Name, proj.ID, proj.BillingID, proj.EnableSandbox)
+	if err != nil {
+		return false, telemetry.Error(ctx, span, err, "error creating Metronome customer")
+	}
+
+	telemetry.WithAttributes(span,
+		telemetry.AttributeKV{Key: "usage-id", Value: proj.UsageID},
+		telemetry.AttributeKV{Key: "usage-plan-id", Value: proj.UsagePlanID},
+	)
+
+	proj.UsageID = customerID
+	proj.UsagePlanID = customerPlanID
+	return true, nil
 }