enkly-ai/services/billing/internal/infrastructure/postgres/subscription_repo.go

// Package postgres implements billing's domain repositories against the
// billing database (see services/billing/migrations/0001_init.sql).
// One-for-one transplant of internal/db/subscriptions.go.
package postgres

import (
	"context"
	"database/sql"
	"errors"
	"time"

	"enkly/pkg/pgkit"
	"enkly/services/billing/internal/domain"
)

const subscriptionColumns = `user_id, plan, status, coalesce(stripe_customer_id, ''), coalesce(stripe_subscription_id, ''), ` +
	`coalesce(stripe_price_id, ''), current_period_end, cancel_at_period_end, created_at, updated_at`

func scanSubscription(row interface{ Scan(...any) error }) (*domain.Subscription, error) {
	var (
		s         domain.Subscription
		periodEnd sql.NullTime
	)
	if err := row.Scan(&s.UserID, &s.Plan, &s.Status, &s.StripeCustomerID, &s.StripeSubscriptionID,
		&s.StripePriceID, &periodEnd, &s.CancelAtPeriodEnd, &s.CreatedAt, &s.UpdatedAt); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, domain.ErrNotFound
		}
		return nil, err
	}
	if periodEnd.Valid {
		s.CurrentPeriodEnd = periodEnd.Time
	}
	return &s, nil
}

type SubscriptionRepo struct {
	DB *pgkit.DB
}

func NewSubscriptionRepo(db *pgkit.DB) *SubscriptionRepo { return &SubscriptionRepo{DB: db} }

func (r *SubscriptionRepo) GetOrCreate(ctx context.Context, userID int64) (*domain.Subscription, error) {
	if _, err := r.DB.ExecContext(ctx,
		`INSERT INTO subscriptions (user_id) VALUES ($1) ON CONFLICT (user_id) DO NOTHING`, userID); err != nil {
		return nil, err
	}
	return scanSubscription(r.DB.QueryRowContext(ctx,
		`SELECT `+subscriptionColumns+` FROM subscriptions WHERE user_id = $1`, userID))
}

func (r *SubscriptionRepo) SetStripeCustomer(ctx context.Context, userID int64, customerID string) error {
	_, err := r.DB.ExecContext(ctx,
		`UPDATE subscriptions SET stripe_customer_id = $2, updated_at = now() WHERE user_id = $1`,
		userID, customerID)
	return err
}

func (r *SubscriptionRepo) ApplyStripeSubscription(ctx context.Context, customerID, subscriptionID, priceID, plan, status string,
	periodEnd time.Time, cancelAtPeriodEnd bool) error {
	var pe any
	if !periodEnd.IsZero() {
		pe = periodEnd
	}
	res, err := r.DB.ExecContext(ctx, `
		UPDATE subscriptions SET
			plan = $2, status = $3, stripe_subscription_id = $4, stripe_price_id = $5,
			current_period_end = $6, cancel_at_period_end = $7, updated_at = now()
		WHERE stripe_customer_id = $1`,
		customerID, plan, status, subscriptionID, priceID, pe, cancelAtPeriodEnd)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return domain.ErrNotFound
	}
	return nil
}

func (r *SubscriptionRepo) ApplyStripeSubscriptionForUser(ctx context.Context, userID int64, customerID, subscriptionID, priceID, plan, status string,
	periodEnd time.Time, cancelAtPeriodEnd bool) error {
	var pe any
	if !periodEnd.IsZero() {
		pe = periodEnd
	}
	_, err := r.DB.ExecContext(ctx, `
		INSERT INTO subscriptions (user_id, plan, status, stripe_customer_id, stripe_subscription_id, stripe_price_id, current_period_end, cancel_at_period_end)
		VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
		ON CONFLICT (user_id) DO UPDATE SET
			plan = excluded.plan, status = excluded.status, stripe_customer_id = excluded.stripe_customer_id,
			stripe_subscription_id = excluded.stripe_subscription_id, stripe_price_id = excluded.stripe_price_id,
			current_period_end = excluded.current_period_end, cancel_at_period_end = excluded.cancel_at_period_end,
			updated_at = now()`,
		userID, plan, status, customerID, subscriptionID, priceID, pe, cancelAtPeriodEnd)
	return err
}

func (r *SubscriptionRepo) CancelByStripeSubscriptionID(ctx context.Context, subscriptionID string) error {
	_, err := r.DB.ExecContext(ctx, `
		UPDATE subscriptions SET
			plan = 'free', status = 'canceled', stripe_subscription_id = NULL,
			stripe_price_id = NULL, cancel_at_period_end = false, updated_at = now()
		WHERE stripe_subscription_id = $1`, subscriptionID)
	return err
}

// ---------------------------------------------------------------------
// Webhook idempotency
// ---------------------------------------------------------------------

type StripeEventRepo struct {
	DB *pgkit.DB
}

func NewStripeEventRepo(db *pgkit.DB) *StripeEventRepo { return &StripeEventRepo{DB: db} }

func (r *StripeEventRepo) RecordIfNew(ctx context.Context, eventID, eventType string) (bool, error) {
	res, err := r.DB.ExecContext(ctx,
		`INSERT INTO stripe_events (id, type) VALUES ($1, $2) ON CONFLICT (id) DO NOTHING`, eventID, eventType)
	if err != nil {
		return false, err
	}
	n, err := res.RowsAffected()
	return n > 0, err
}

// ---------------------------------------------------------------------
// Daily usage
// ---------------------------------------------------------------------

type UsageRepo struct {
	DB *pgkit.DB
}

func NewUsageRepo(db *pgkit.DB) *UsageRepo { return &UsageRepo{DB: db} }

func (r *UsageRepo) DailyUsage(ctx context.Context, userID int64) (int, error) {
	var n int
	err := r.DB.QueryRowContext(ctx,
		`SELECT messages FROM usage_daily WHERE user_id = $1 AND day = current_date`, userID).Scan(&n)
	if errors.Is(err, sql.ErrNoRows) {
		return 0, nil
	}
	return n, err
}

func (r *UsageRepo) IncrementDailyUsage(ctx context.Context, userID int64) (int, error) {
	var n int
	err := r.DB.QueryRowContext(ctx, `
		INSERT INTO usage_daily (user_id, day, messages) VALUES ($1, current_date, 1)
		ON CONFLICT (user_id, day) DO UPDATE SET messages = usage_daily.messages + 1
		RETURNING messages`, userID).Scan(&n)
	return n, err
}