geo/server/internal/tenant/repository/quota_repo.go

package repository

import (
	"context"
	"time"

	"github.com/geo-platform/tenant-api/internal/tenant/repository/generated"
)

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)
	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
}

func NewQuotaRepository(db generated.DBTX) QuotaRepository {
	return &quotaRepository{q: newQuerier(db)}
}

func (r *quotaRepository) GetCurrentBalance(ctx context.Context, tenantID int64, quotaType string) (int, error) {
	balance, err := r.q.GetCurrentBalance(ctx, generated.GetCurrentBalanceParams{
		TenantID:  tenantID,
		QuotaType: quotaType,
	})
	if err != nil {
		return 0, err
	}
	return int(balance), 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,
	})
}