94 lines
3.3 KiB
Python
94 lines
3.3 KiB
Python
from decimal import Decimal
|
|
|
|
from django.db import transaction
|
|
|
|
from apps.billing.models import CreditAccount, CreditLedger, CreditReservation
|
|
|
|
|
|
@transaction.atomic
|
|
def reserve_credit(*, team, user, task, amount: Decimal) -> CreditReservation:
|
|
account, _ = CreditAccount.objects.select_for_update().get_or_create(team=team)
|
|
available = account.balance - account.reserved_balance
|
|
if available < amount:
|
|
raise ValueError("insufficient credit")
|
|
|
|
account.reserved_balance += amount
|
|
account.save(update_fields=["reserved_balance", "updated_at"])
|
|
reservation = CreditReservation.objects.create(
|
|
team=team,
|
|
user=user,
|
|
project=task.project,
|
|
task=task,
|
|
amount=amount,
|
|
)
|
|
CreditLedger.objects.create(
|
|
team=team,
|
|
user=user,
|
|
project=task.project,
|
|
task=task,
|
|
ledger_type=CreditLedger.Type.RESERVE,
|
|
amount=amount,
|
|
balance_after=account.balance,
|
|
reason="reserve ai task credit",
|
|
)
|
|
return reservation
|
|
|
|
|
|
@transaction.atomic
|
|
def release_credit(*, reservation: CreditReservation, reason: str = "") -> None:
|
|
account = CreditAccount.objects.select_for_update().get(team=reservation.team)
|
|
if reservation.status != CreditReservation.Status.ACTIVE:
|
|
return
|
|
|
|
account.reserved_balance -= reservation.amount
|
|
account.save(update_fields=["reserved_balance", "updated_at"])
|
|
reservation.status = CreditReservation.Status.RELEASED
|
|
reservation.save(update_fields=["status", "updated_at"])
|
|
CreditLedger.objects.create(
|
|
team=reservation.team,
|
|
user=reservation.user,
|
|
project=reservation.project,
|
|
task=reservation.task,
|
|
ledger_type=CreditLedger.Type.RELEASE,
|
|
amount=reservation.amount,
|
|
balance_after=account.balance,
|
|
reason=reason or "release reserved credit",
|
|
)
|
|
|
|
|
|
@transaction.atomic
|
|
def charge_reserved_credit(*, reservation: CreditReservation, actual_amount: Decimal) -> None:
|
|
account = CreditAccount.objects.select_for_update().get(team=reservation.team)
|
|
if reservation.status != CreditReservation.Status.ACTIVE:
|
|
raise ValueError("reservation is not active")
|
|
if actual_amount > reservation.amount:
|
|
raise ValueError("actual amount exceeds reserved amount")
|
|
|
|
account.balance -= actual_amount
|
|
account.reserved_balance -= reservation.amount
|
|
account.save(update_fields=["balance", "reserved_balance", "updated_at"])
|
|
reservation.status = CreditReservation.Status.CHARGED
|
|
reservation.save(update_fields=["status", "updated_at"])
|
|
CreditLedger.objects.create(
|
|
team=reservation.team,
|
|
user=reservation.user,
|
|
project=reservation.project,
|
|
task=reservation.task,
|
|
ledger_type=CreditLedger.Type.CHARGE,
|
|
amount=actual_amount,
|
|
balance_after=account.balance,
|
|
reason="charge ai task credit",
|
|
)
|
|
if reservation.amount > actual_amount:
|
|
CreditLedger.objects.create(
|
|
team=reservation.team,
|
|
user=reservation.user,
|
|
project=reservation.project,
|
|
task=reservation.task,
|
|
ledger_type=CreditLedger.Type.RELEASE,
|
|
amount=reservation.amount - actual_amount,
|
|
balance_after=account.balance,
|
|
reason="release unused reserved credit",
|
|
)
|
|
|