billing errors and warnings, stripe and orb webhook, river worker for async background jobs

admin/billing.go

@@ -4,13 +4,15 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"time"
 
 	"github.com/rilldata/rill/admin/billing"
 	"github.com/rilldata/rill/admin/database"
 	"go.uber.org/zap"
 )
 
 func (s *Service) InitOrganizationBilling(ctx context.Context, org *database.Organization) (*database.Organization, *billing.Subscription, error) {
+	// TODO This can be moved to a background job and repair org billing job can be removed in the next version. We need repair job to fix existing orgs but afterwards background job wil ensure that all orgs are in sync with billing system
 	// create payment customer
 	pc, err := s.PaymentProvider.CreateCustomer(ctx, org)
 	if err != nil {
@@ -39,6 +41,34 @@ func (s *Service) InitOrganizationBilling(ctx context.Context, org *database.Org
 	}
 	s.Logger.Info("created subscription", zap.String("org", org.Name), zap.String("subscription_id", sub.ID))
 
+	// scheduling it before the org update as repair billing job can take care of it if update fails
+	err = s.ScheduleTrialEndCheckJobs(ctx, org.ID, sub.ID, plan.ID, sub.StartDate, sub.TrialEndDate)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	// raise no payment method billing error
+	_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+		OrgID:     org.ID,
+		Type:      database.BillingErrorTypeNoPaymentMethod,
+		Metadata:  &database.BillingErrorMetadataNoPaymentMethod{},
+		EventTime: org.CreatedOn,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+	}
+
+	// raise no billable address billing error
+	_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+		OrgID:     org.ID,
+		Type:      database.BillingErrorTypeNoBillableAddress,
+		Metadata:  &database.BillingErrorMetadataNoBillableAddress{},
+		EventTime: org.CreatedOn,
+	})
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+	}
+
 	org, err = s.DB.UpdateOrganization(ctx, org.ID, &database.UpdateOrganizationOptions{
 		Name:                                org.Name,
 		DisplayName:                         org.DisplayName,
@@ -81,6 +111,34 @@ func (s *Service) RepairOrgBilling(ctx context.Context, org *database.Organizati
 			}
 			s.Logger.Info("created subscription", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("subscription_id", sub.ID))
 			subs = append(subs, sub)
+
+			// schedule trial end check job to the river queue, if it was already scheduled it will be ignored
+			err = s.ScheduleTrialEndCheckJobs(ctx, org.ID, sub.ID, plan.ID, sub.StartDate, sub.TrialEndDate)
+			if err != nil {
+				return nil, nil, err
+			}
+
+			// raise no payment method billing error
+			_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+				OrgID:     org.ID,
+				Type:      database.BillingErrorTypeNoPaymentMethod,
+				Metadata:  &database.BillingErrorMetadataNoPaymentMethod{},
+				EventTime: org.CreatedOn,
+			})
+			if err != nil {
+				return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+			}
+
+			// raise no billable address billing error
+			_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+				OrgID:     org.ID,
+				Type:      database.BillingErrorTypeNoBillableAddress,
+				Metadata:  &database.BillingErrorMetadataNoBillableAddress{},
+				EventTime: org.CreatedOn,
+			})
+			if err != nil {
+				return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+			}
 		}
 		if len(subs) > 1 {
 			s.Logger.Warn("multiple subscriptions found for the customer", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Int("num_subscriptions", len(subs)))
@@ -152,6 +210,34 @@ func (s *Service) RepairOrgBilling(ctx context.Context, org *database.Organizati
 		}
 		s.Logger.Info("created subscription", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("subscription_id", sub.ID))
 		subs = append(subs, sub)
+
+		// schedule trial end check job to the river queue
+		err = s.ScheduleTrialEndCheckJobs(ctx, org.ID, sub.ID, plan.ID, sub.StartDate, sub.TrialEndDate)
+		if err != nil {
+			return nil, nil, err
+		}
+
+		// raise no payment method billing error
+		_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+			OrgID:     org.ID,
+			Type:      database.BillingErrorTypeNoPaymentMethod,
+			Metadata:  &database.BillingErrorMetadataNoPaymentMethod{},
+			EventTime: org.CreatedOn,
+		})
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+		}
+
+		// raise no billable address billing error
+		_, err = s.DB.UpsertBillingError(ctx, &database.UpsertBillingErrorOptions{
+			OrgID:     org.ID,
+			Type:      database.BillingErrorTypeNoBillableAddress,
+			Metadata:  &database.BillingErrorMetadataNoBillableAddress{},
+			EventTime: org.CreatedOn,
+		})
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to upsert billing error: %w", err)
+		}
 	} else if len(subs) > 1 {
 		s.Logger.Warn("multiple subscriptions found for the customer", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Int("num_subscriptions", len(subs)))
 	}
@@ -179,6 +265,151 @@ func (s *Service) RepairOrgBilling(ctx context.Context, org *database.Organizati
 	return org, subs, nil
 }
 
+func (s *Service) ScheduleTrialEndCheckJobs(ctx context.Context, orgID, subID, planID string, trialStartDate, trialEndDate time.Time) error {
+	if trialEndDate.Before(time.Now()) {
+		return nil
+	}
+
+	// schedule trial ending soon job
+	tes, err := s.Jobs.TrialEndingSoon(ctx, orgID, subID, planID, trialEndDate)
+	if err != nil {
+		return fmt.Errorf("failed to schedule trial ending soon job: %w", err)
+	}
+
+	// schedule trial end check job
+	tec, err := s.Jobs.TrialEndCheck(ctx, orgID, subID, planID, trialEndDate)
+	if err != nil {
+		return fmt.Errorf("failed to schedule trial end check job: %w", err)
+	}
+
+	// raise on-trial billing warning
+	_, err = s.DB.UpsertBillingWarning(ctx, &database.UpsertBillingWarningOptions{
+		OrgID: orgID,
+		Type:  database.BillingWarningTypeOnTrial,
+		Metadata: &database.BillingWarningMetadataOnTrial{
+			EndDate:              trialEndDate,
+			TrialEndingSoonJobID: tes.ID,
+			TrialEndCheckJobID:   tec.ID,
+		},
+		EventTime: trialStartDate,
+	})
+	if err != nil {
+		return fmt.Errorf("failed to upsert billing warning: %w", err)
+	}
+
+	return nil
+}
+
+// CleanupTrialBillingErrorsAndWarnings removes trial related billing error and warning and cancel associated jobs
+func (s *Service) CleanupTrialBillingErrorsAndWarnings(ctx context.Context, orgID string) error {
+	bett, err := s.DB.FindBillingErrorByType(ctx, orgID, database.BillingErrorTypeTrialEnded)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return fmt.Errorf("failed to find billing errors: %w", err)
+		}
+	}
+
+	if bett != nil {
+		metadata, ok := bett.Metadata.(*database.BillingErrorMetadataTrialEnded)
+		if ok && metadata.GracePeriodEndJobID > 0 {
+			// cancel the trial end grace check job, ignore errors.
+			_ = s.Jobs.CancelJob(ctx, metadata.GracePeriodEndJobID)
+		}
+		err = s.DB.DeleteBillingError(ctx, bett.ID)
+		if err != nil {
+			return fmt.Errorf("failed to delete billing error: %w", err)
+		}
+	}
+
+	bwot, err := s.DB.FindBillingWarningByType(ctx, orgID, database.BillingWarningTypeOnTrial)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return fmt.Errorf("failed to find billing warnings: %w", err)
+		}
+	}
+
+	if bwot != nil {
+		metadata, ok := bwot.Metadata.(*database.BillingWarningMetadataOnTrial)
+		if ok && metadata.TrialEndingSoonJobID > 0 {
+			// cancel the trial ending soon job, ignore errors.
+			_ = s.Jobs.CancelJob(ctx, metadata.TrialEndingSoonJobID)
+		}
+
+		if ok && metadata.TrialEndCheckJobID > 0 {
+			// cancel the trial end check job, ignore errors.
+			_ = s.Jobs.CancelJob(ctx, metadata.TrialEndCheckJobID)
+		}
+
+		err = s.DB.DeleteBillingWarning(ctx, bwot.ID)
+		if err != nil {
+			return fmt.Errorf("failed to delete billing warning: %w", err)
+		}
+	}
+
+	return nil
+}
+
+// CleanupBillingErrorSubCancellation removes subscription cancellation related billing error and cancel associated job
+func (s *Service) CleanupBillingErrorSubCancellation(ctx context.Context, orgID string) error {
+	besc, err := s.DB.FindBillingErrorByType(ctx, orgID, database.BillingErrorTypeSubscriptionCancelled)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return fmt.Errorf("failed to find billing errors: %w", err)
+		}
+	}
+
+	if besc != nil {
+		metadata, ok := besc.Metadata.(*database.BillingErrorMetadataSubscriptionCancelled)
+		if ok && metadata.SubEndJobID > 0 {
+			// cancel the subscription end check job, ignore errors.
+			_ = s.Jobs.CancelJob(ctx, metadata.SubEndJobID)
+		}
+		err = s.DB.DeleteBillingError(ctx, besc.ID)
+		if err != nil {
+			return fmt.Errorf("failed to delete billing error: %w", err)
+		}
+	}
+
+	return nil
+}
+
+func (s *Service) CheckBillingErrors(ctx context.Context, orgID string) error {
+	be, err := s.DB.FindBillingErrorByType(ctx, orgID, database.BillingErrorTypeTrialEnded)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return err
+		}
+	}
+
+	if be != nil {
+		return fmt.Errorf("trial has ended")
+	}
+
+	be, err = s.DB.FindBillingErrorByType(ctx, orgID, database.BillingErrorTypeInvoicePaymentFailed)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return err
+		}
+	}
+
+	if be != nil { // should we allow any grace period here?
+		return fmt.Errorf("invoice payment failed")
+	}
+
+	be, err = s.DB.FindBillingErrorByType(ctx, orgID, database.BillingErrorTypeSubscriptionCancelled)
+	if err != nil {
+		if !errors.Is(err, database.ErrNotFound) {
+			return err
+		}
+	}
+
+	if be != nil && be.Metadata.(*database.BillingErrorMetadataSubscriptionCancelled).EndDate.AddDate(0, 0, 1).After(time.Now()) {
+		return fmt.Errorf("subscription cancelled")
+	}
+
+	return nil
+}
+
 func valOrDefault[T any](ptr *T, def T) T {
 	if ptr != nil {
 		return *ptr

admin/billing/biller.go

@@ -3,9 +3,11 @@ package billing
 import (
 	"context"
 	"errors"
+	"net/http"
 	"time"
 
 	"github.com/rilldata/rill/admin/database"
+	"github.com/rilldata/rill/admin/jobs"
 )
 
 const (
@@ -32,21 +34,32 @@ type Biller interface {
 	UpdateCustomerPaymentID(ctx context.Context, customerID string, provider PaymentProvider, paymentProviderID string) error
 	UpdateCustomerEmail(ctx context.Context, customerID, email string) error
 
-	// CreateSubscription creates a subscription for the given organization.
-	// The subscription starts immediately.
+	// CreateSubscription creates a subscription for the given organization. Subscription starts immediately.
 	CreateSubscription(ctx context.Context, customerID string, plan *Plan) (*Subscription, error)
+	// CreateSubscriptionInFuture creates a subscription for the given organization with a start date in the future.
+	CreateSubscriptionInFuture(ctx context.Context, customerID string, plan *Plan, startDate time.Time) (*Subscription, error)
 	CancelSubscription(ctx context.Context, subscriptionID string, cancelOption SubscriptionCancellationOption) error
 	GetSubscriptionsForCustomer(ctx context.Context, customerID string) ([]*Subscription, error)
-	ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan) (*Subscription, error)
+	ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan, changeOption SubscriptionChangeOption) (*Subscription, error)
 	// CancelSubscriptionsForCustomer deletes the subscription for the given organization.
 	// cancellationDate only applicable if option is SubscriptionCancellationOptionRequestedDate
 	CancelSubscriptionsForCustomer(ctx context.Context, customerID string, cancelOption SubscriptionCancellationOption) error
 	FindSubscriptionsPastTrialPeriod(ctx context.Context) ([]*Subscription, error)
 
+	GetInvoice(ctx context.Context, invoiceID string) (*Invoice, error)
+	IsInvoiceValid(ctx context.Context, invoice *Invoice) bool
+	IsInvoicePaid(ctx context.Context, invoice *Invoice) bool
+
 	ReportUsage(ctx context.Context, usage []*Usage) error
 
 	GetReportingGranularity() UsageReportingGranularity
 	GetReportingWorkerCron() string
+
+	// WebhookHandlerFunc returns a http.HandlerFunc that can be used to handle incoming webhooks from the payment provider. Return nil if you don't want to register any webhook handlers.
+	WebhookHandlerFunc(ctx context.Context) func(w http.ResponseWriter, r *http.Request)
+
+	// SetJobsClient needs to be explicitly set because of circular dependency, see admin start method
+	SetJobsClient(jobs jobs.Client)
 }
 
 type Plan struct {
@@ -96,11 +109,12 @@ type Subscription struct {
 }
 
 type Customer struct {
-	ID                string
-	Email             string
-	Name              string
-	PaymentProviderID string
-	PortalURL         string
+	ID                 string
+	Email              string
+	Name               string
+	PaymentProviderID  string
+	PortalURL          string
+	HasBillableAddress bool
 }
 
 type Usage struct {
@@ -113,6 +127,18 @@ type Usage struct {
 	Metadata       map[string]interface{}
 }
 
+type Invoice struct {
+	ID             string
+	Status         string
+	CustomerID     string
+	Amount         string
+	Currency       string
+	DueDate        time.Time
+	CreatedAt      time.Time
+	SubscriptionID string
+	Metadata       map[string]interface{}
+}
+
 type UsageReportingGranularity string
 
 const (
@@ -127,6 +153,13 @@ const (
 	SubscriptionCancellationOptionImmediate
 )
 
+type SubscriptionChangeOption int
+
+const (
+	SubscriptionChangeOptionEndOfSubscriptionTerm SubscriptionChangeOption = iota
+	SubscriptionChangeOptionImmediate
+)
+
 type PaymentProvider string
 
 const (