Improve billing and usage customer creation, handle errors for admin user

api/server/handlers/billing/list.go

@@ -1,6 +1,7 @@
 package billing
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 
@@ -65,24 +66,94 @@ 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
+	}
+
+	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: "payment-enabled", Value: paymentEnabled},
+	)
+
+	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()
+
+	if proj.BillingID == "" || proj.UsageID == uuid.Nil {
+		adminUser, err := c.getAdminUser(ctx, proj.ID)
+		if err != nil {
+			err = telemetry.Error(ctx, span, err, "error getting admin user")
+			return err
+		}
+
+		// If the admin user is not found, use the current user as last resort
+		if adminUser == nil {
+			adminUser = user
+		}
+
+		// 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 ensuring Stripe customer exists")
+			return err
+		}
+
+		// 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 {
+			err = telemetry.Error(ctx, span, err, "error ensuring Metronome customer exists")
+			return err
+		}
+
+		if shouldUpdateBilling || shouldUpdateUsage {
+			_, err := c.Repo().Project().UpdateProject(proj)
+			if err != nil {
+				err = telemetry.Error(ctx, span, err, "error updating project")
+				return err
+			}
+		}
+	}
+
+	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 {
+		err = telemetry.Error(ctx, span, err, "error listing project roles")
+		return adminUser, err
+	}
+
 	// 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
+				// If the user is not found, continue to the next role
+				continue
 			}
 			break
 		}
@@ -93,59 +164,55 @@ func (c *CheckPaymentEnabledHandler) ServeHTTP(w http.ResponseWriter, r *http.Re
 		telemetry.AttributeKV{Key: "admin-user-email", Value: adminUser.Email},
 	)
 
-	// 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)
-		if err != nil {
-			err = telemetry.Error(ctx, span, err, "error creating billing customer")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
-			return
-		}
-		proj.BillingID = billingID
-		shouldUpdate = true
+	return adminUser, nil
+}
 
-		telemetry.WithAttributes(span,
-			telemetry.AttributeKV{Key: "billing-id", Value: proj.BillingID},
-		)
+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 creating billing customer")
+		return false, err
+	}
+
+	telemetry.WithAttributes(span,
+		telemetry.AttributeKV{Key: "billing-id", Value: proj.BillingID},
+	)
+
+	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 {
-		customerID, customerPlanID, err := c.Config().BillingManager.MetronomeClient.CreateCustomerWithPlan(ctx, adminUser.Email, proj.Name, proj.ID, proj.BillingID, proj.EnableSandbox)
+		customerID, customerPlanID, err := c.Config().BillingManager.MetronomeClient.CreateCustomerWithPlan(ctx, adminUserEmail, proj.Name, proj.ID, proj.BillingID, proj.EnableSandbox)
 		if err != nil {
 			err = telemetry.Error(ctx, span, err, "error creating Metronome customer")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
+			return false, err
 		}
-		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)
-		if err != nil {
-			err := telemetry.Error(ctx, span, err, "error updating project")
-			c.HandleAPIError(w, r, apierrors.NewErrInternal(err))
-			return
-		}
-	}
-
-	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(fmt.Errorf("error checking if payment enabled: %w", err)))
-		return
+		proj.UsageID = customerID
+		proj.UsagePlanID = customerPlanID
+		return true, nil
 	}
 
-	telemetry.WithAttributes(span,
-		telemetry.AttributeKV{Key: "project-id", Value: proj.ID},
-		telemetry.AttributeKV{Key: "customer-id", Value: proj.BillingID},
-	)
-
-	c.WriteResult(w, r, paymentEnabled)
+	return false, nil
 }