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)