thread-master/apps/backend/internal/module/billing/mongo_repository.go

package billing

import (
	"context"
	"time"

	libmongo "apps/backend/internal/lib/mongo"
	usageDomain "apps/backend/internal/module/usage/domain"

	"github.com/zeromicro/go-zero/core/stores/mon"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

type MongoRepository struct {
	attempts  *mon.Model
	prefs     *mon.Model
	purchases *mon.Model
	webhooks  *mon.Model
}

func NewMongoRepository(uri, database string) *MongoRepository {
	uri = libmongo.MustMongoURI(uri)
	return &MongoRepository{
		attempts:  mon.MustNewModel(uri, database, "billing_checkout_attempts"),
		prefs:     mon.MustNewModel(uri, database, "usage_prefs"),
		purchases: mon.MustNewModel(uri, database, "usage_purchases"),
		webhooks:  mon.MustNewModel(uri, database, "billing_webhook_events"),
	}
}

func (r *MongoRepository) InsertCheckout(ctx context.Context, a *CheckoutAttempt) (bool, error) {
	_, err := r.attempts.InsertOne(ctx, a)
	if mongo.IsDuplicateKeyError(err) {
		return false, nil
	}
	return err == nil, err
}

func (r *MongoRepository) GetCheckoutByRequest(ctx context.Context, uid int64, requestID string) (*CheckoutAttempt, error) {
	return r.findAttempt(ctx, bson.M{"uid": uid, "request_id": requestID})
}
func (r *MongoRepository) GetCheckout(ctx context.Context, uid int64, id string) (*CheckoutAttempt, error) {
	return r.findAttempt(ctx, bson.M{"_id": id, "uid": uid})
}
func (r *MongoRepository) GetCheckoutBySession(ctx context.Context, id string) (*CheckoutAttempt, error) {
	return r.findAttempt(ctx, bson.M{"session_id": id})
}
func (r *MongoRepository) findAttempt(ctx context.Context, filter bson.M) (*CheckoutAttempt, error) {
	var a CheckoutAttempt
	if err := r.attempts.FindOne(ctx, &a, filter); err != nil {
		if err == mon.ErrNotFound {
			return nil, ErrNotFound
		}
		return nil, err
	}
	return &a, nil
}
func (r *MongoRepository) UpdateCheckout(ctx context.Context, a *CheckoutAttempt) error {
	_, err := r.attempts.ReplaceOne(ctx, bson.M{"_id": a.ID}, a)
	return err
}

func (r *MongoRepository) GetSubscription(ctx context.Context, uid int64) (*SubscriptionState, error) {
	var p usageDomain.MemberPrefs
	if err := r.prefs.FindOne(ctx, &p, bson.M{"uid": uid}); err != nil {
		if err == mon.ErrNotFound {
			return &SubscriptionState{UID: uid, PlanID: PlanFree, Status: "none"}, nil
		}
		return nil, err
	}
	return &SubscriptionState{UID: uid, PlanID: p.PlanID, PaidPlanID: p.PaidPlanID, Status: p.SubscriptionStatus, CustomerID: p.StripeCustomerID, SubscriptionID: p.StripeSubscriptionID, CurrentPeriodStart: p.CurrentPeriodStart, CurrentPeriodEnd: p.CurrentPeriodEnd, CancelAtPeriodEnd: p.CancelAtPeriodEnd, AdminOverride: p.AdminOverride}, nil
}

func (r *MongoRepository) FindUIDByCustomer(ctx context.Context, customerID string) (int64, error) {
	var p usageDomain.MemberPrefs
	if customerID == "" {
		return 0, ErrNotFound
	}
	if err := r.prefs.FindOne(ctx, &p, bson.M{"stripe_customer_id": customerID}); err != nil {
		if err == mon.ErrNotFound {
			a, attemptErr := r.findAttempt(ctx, bson.M{"customer_id": customerID})
			if attemptErr != nil {
				return 0, ErrNotFound
			}
			return a.UID, nil
		}
		return 0, err
	}
	return p.UID, nil
}

func (r *MongoRepository) ApplyEntitlement(ctx context.Context, u EntitlementUpdate) (bool, error) {
	filter := bson.M{"uid": u.UID, "$or": bson.A{
		bson.M{"billing_event_created_at": bson.M{"$exists": false}},
		bson.M{"billing_event_created_at": bson.M{"$lt": u.EventCreatedAt}},
		bson.M{"billing_event_created_at": u.EventCreatedAt, "billing_event_id": bson.M{"$lt": u.EventID}},
	}}
	pipeline := mongo.Pipeline{bson.D{{Key: "$set", Value: bson.D{
		{Key: "uid", Value: u.UID},
		{Key: "paid_plan_id", Value: u.PaidPlanID},
		{Key: "subscription_status", Value: u.Status},
		{Key: "stripe_customer_id", Value: u.CustomerID},
		{Key: "stripe_subscription_id", Value: u.SubscriptionID},
		{Key: "current_period_start", Value: u.CurrentPeriodStart},
		{Key: "current_period_end", Value: u.CurrentPeriodEnd},
		{Key: "cancel_at_period_end", Value: u.CancelAtPeriodEnd},
		{Key: "billing_event_id", Value: u.EventID},
		{Key: "billing_event_created_at", Value: u.EventCreatedAt},
		{Key: "updated_at", Value: time.Now().UTC().UnixNano()},
		{Key: "plan_id", Value: bson.D{{Key: "$cond", Value: bson.A{
			bson.D{{Key: "$eq", Value: bson.A{bson.D{{Key: "$ifNull", Value: bson.A{"$admin_override", false}}}, true}}},
			bson.D{{Key: "$ifNull", Value: bson.A{"$plan_id", PlanFree}}}, u.PlanID,
		}}}},
	}}}}
	res, err := r.prefs.UpdateOne(ctx, filter, pipeline, options.Update().SetUpsert(true))
	if err != nil {
		if mongo.IsDuplicateKeyError(err) {
			return false, nil
		}
		return false, err
	}
	return res.MatchedCount > 0 || res.UpsertedCount > 0, nil
}

func (r *MongoRepository) WebhookProcessed(ctx context.Context, id string) (bool, error) {
	var out bson.M
	err := r.webhooks.FindOne(ctx, &out, bson.M{"_id": id})
	if err == mon.ErrNotFound {
		return false, nil
	}
	return err == nil, err
}
func (r *MongoRepository) MarkWebhookProcessed(ctx context.Context, id, eventType string, created int64) error {
	_, err := r.webhooks.ReplaceOne(ctx, bson.M{"_id": id}, bson.M{"_id": id, "stripe_event_id": id, "type": eventType, "stripe_created_at": created * int64(time.Second), "processed_at": time.Now().UTC().UnixNano()}, options.Replace().SetUpsert(true))
	return err
}
func (r *MongoRepository) UpsertPurchase(ctx context.Context, p *PurchaseAudit) error {
	_, err := r.purchases.ReplaceOne(ctx, bson.M{"stripe_event_id": p.StripeEventID}, p, options.Replace().SetUpsert(true))
	return err
}

var _ Repository = (*MongoRepository)(nil)