package repository import ( "context" "errors" "time" "github.com/jackc/pgx/v5" "github.com/geo-platform/tenant-api/internal/tenant/repository/generated" ) type ArticleQuotaStatus struct { PlanCode string PlanName string Total int Used int Balance int ResetAt *time.Time } type AIQuotaStatus struct { PlanCode string PlanName string Total int Used int Balance int BaseChars int ResetAt *time.Time } type QuotaLedgerInput struct { TenantID int64 OperatorID int64 QuotaType string Delta int BalanceAfter int Reason *string ReferenceType *string ReferenceID *int64 } type QuotaReservationInput struct { TenantID int64 OperatorID int64 QuotaType string ResourceType string ResourceID *int64 ReservedAmount int ExpireAt time.Time } type QuotaRepository interface { GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error RefundReservation(ctx context.Context, reservationID, tenantID int64) error } type quotaRepository struct { q generated.Querier db generated.DBTX } func NewQuotaRepository(db generated.DBTX) QuotaRepository { return "aRepository{q: newQuerier(db), db: db} } func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error) { if quotaType != "article_generation" { balance, err := r.q.GetCurrentBalance(ctx, generated.GetCurrentBalanceParams{ TenantID: tenantID, QuotaType: quotaType, }) if err != nil { return 0, err } return int(balance), nil } status, err := r.GetArticleQuotaStatus(ctx, tenantID, time.Now().UTC()) if err != nil { return 0, err } return status.Balance, nil } func (r *quotaRepository) GetAIQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*AIQuotaStatus, error) { access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now) if err != nil { return nil, err } if access == nil { return nil, pgx.ErrNoRows } baseChars := access.AIPointBaseChars() if !access.HasActiveAccess(now) { return &AIQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, BaseChars: baseChars, }, nil } total := access.AIPointsMonthlyLimit() if total <= 0 { return &AIQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, BaseChars: baseChars, }, nil } windowStart, windowEnd, resetAt := access.AIPointsWindow(now) if windowEnd.IsZero() { return nil, errors.New("ai points quota window end is missing") } used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "ai_points", windowStart, windowEnd) if err != nil { return nil, err } balance := total - used if balance < 0 { balance = 0 } return &AIQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, Total: total, Used: used, Balance: balance, BaseChars: baseChars, ResetAt: resetAt, }, nil } func (r *quotaRepository) GetArticleQuotaStatus(ctx context.Context, tenantID int64, now time.Time) (*ArticleQuotaStatus, error) { access, err := loadTenantPlanAccess(ctx, r.db, tenantID, now) if err != nil { return nil, err } if access == nil { return nil, pgx.ErrNoRows } // Non-HTTP callers (scheduler, workers) bypass subscription_guard, so any // code path that consults article quota must itself enforce that the // subscription is currently active. Return a zero-balance status when the // plan is expired/inactive/not-yet-active so the caller short-circuits on // insufficient balance instead of dispatching generation. if !access.HasActiveAccess(now) { return &ArticleQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, Total: 0, Used: 0, Balance: 0, }, nil } total := access.ArticleGenerationLimit() if total <= 0 { return &ArticleQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, Total: 0, Used: 0, Balance: 0, }, nil } windowStart, windowEnd, resetAt := access.ArticleQuotaWindow(now) if windowEnd.IsZero() { return nil, errors.New("article quota window end is missing") } used, err := CountEffectiveQuotaReservations(ctx, r.db, tenantID, "article_generation", windowStart, windowEnd) if err != nil { return nil, err } balance := total - used if balance < 0 { balance = 0 } return &ArticleQuotaStatus{ PlanCode: access.PlanCode, PlanName: access.PlanName, Total: total, Used: used, Balance: balance, ResetAt: resetAt, }, nil } func (r *quotaRepository) InsertQuotaLedger(ctx context.Context, input QuotaLedgerInput) (int64, error) { var operatorID *int64 if input.OperatorID > 0 { operatorID = &input.OperatorID } return r.q.InsertQuotaLedger(ctx, generated.InsertQuotaLedgerParams{ TenantID: input.TenantID, OperatorID: pgInt8(operatorID), QuotaType: input.QuotaType, Delta: int32(input.Delta), BalanceAfter: int32(input.BalanceAfter), Reason: pgText(input.Reason), ReferenceType: pgText(input.ReferenceType), ReferenceID: pgInt8(input.ReferenceID), }) } func (r *quotaRepository) CreateQuotaReservation(ctx context.Context, input QuotaReservationInput) (int64, error) { expireAt := input.ExpireAt var operatorID *int64 if input.OperatorID > 0 { operatorID = &input.OperatorID } return r.q.CreateQuotaReservation(ctx, generated.CreateQuotaReservationParams{ TenantID: input.TenantID, OperatorID: pgInt8(operatorID), QuotaType: input.QuotaType, ResourceType: input.ResourceType, ResourceID: pgInt8(input.ResourceID), ReservedAmount: int32(input.ReservedAmount), ExpireAt: pgTimestamp(&expireAt), }) } func (r *quotaRepository) UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error { resourceIDValue := resourceID return r.q.UpdateQuotaReservationResource(ctx, generated.UpdateQuotaReservationResourceParams{ ResourceID: pgInt8(&resourceIDValue), ID: reservationID, TenantID: tenantID, }) } func (r *quotaRepository) ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error { return r.q.ConfirmReservation(ctx, generated.ConfirmReservationParams{ ID: reservationID, TenantID: tenantID, }) } func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID, tenantID int64) error { return r.q.RefundReservation(ctx, generated.RefundReservationParams{ ID: reservationID, TenantID: tenantID, }) }