200 lines
6.6 KiB
Python
200 lines
6.6 KiB
Python
"""Wallet service: atomic credit/debit, row locking, and idempotency.
|
|
|
|
docs/planning/economy.md US-1 (task 909e3e9a59af).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
from sqlalchemy import func, select
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.wallet import Transaction, TransactionType, Wallet
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class InsufficientFundsError(Exception):
|
|
"""Raised when a debit would cause wallet balance to drop below zero."""
|
|
pass
|
|
|
|
|
|
class InvalidAmountError(ValueError):
|
|
"""Raised when an invalid (non-positive) amount is supplied."""
|
|
pass
|
|
|
|
|
|
class WalletService:
|
|
@staticmethod
|
|
async def get_or_create_wallet(
|
|
db: AsyncSession, user_id: int, for_update: bool = False
|
|
) -> Wallet:
|
|
"""Fetch existing wallet or lazily create one for user_id.
|
|
|
|
If for_update is True, locks the row with SELECT ... FOR UPDATE.
|
|
"""
|
|
stmt = select(Wallet).where(Wallet.user_id == user_id)
|
|
if for_update:
|
|
stmt = stmt.with_for_update()
|
|
|
|
result = await db.execute(stmt)
|
|
wallet = result.scalar_one_or_none()
|
|
if wallet is not None:
|
|
return wallet
|
|
|
|
# If not found and not in update mode, create lazily
|
|
try:
|
|
new_wallet = Wallet(user_id=user_id, balance=0)
|
|
db.add(new_wallet)
|
|
await db.commit()
|
|
await db.refresh(new_wallet)
|
|
|
|
if for_update:
|
|
stmt_locked = select(Wallet).where(Wallet.id == new_wallet.id).with_for_update()
|
|
res = await db.execute(stmt_locked)
|
|
return res.scalar_one()
|
|
|
|
return new_wallet
|
|
except IntegrityError:
|
|
# Concurrent insert collision -- rollback and fetch the existing one
|
|
await db.rollback()
|
|
stmt = select(Wallet).where(Wallet.user_id == user_id)
|
|
if for_update:
|
|
stmt = stmt.with_for_update()
|
|
result = await db.execute(stmt)
|
|
return result.scalar_one()
|
|
|
|
@classmethod
|
|
async def credit(
|
|
cls,
|
|
db: AsyncSession,
|
|
user_id: int,
|
|
amount: int,
|
|
tx_type: TransactionType | str,
|
|
reference_type: str | None = None,
|
|
reference_id: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> tuple[Wallet, Transaction]:
|
|
"""Credit a user's wallet atomically with row locking.
|
|
|
|
If an idempotency_key is provided and already exists, returns the original
|
|
outcome without double-crediting.
|
|
"""
|
|
if amount <= 0:
|
|
raise InvalidAmountError("Credit amount must be strictly positive")
|
|
|
|
if isinstance(tx_type, str):
|
|
tx_type = TransactionType(tx_type)
|
|
|
|
# Check idempotency first
|
|
if idempotency_key:
|
|
stmt = select(Transaction).where(Transaction.idempotency_key == idempotency_key)
|
|
existing_tx = (await db.execute(stmt)).scalar_one_or_none()
|
|
if existing_tx is not None:
|
|
wallet = await cls.get_or_create_wallet(db, user_id, for_update=False)
|
|
return wallet, existing_tx
|
|
|
|
# Lock wallet row
|
|
wallet = await cls.get_or_create_wallet(db, user_id, for_update=True)
|
|
wallet.balance += amount
|
|
|
|
tx = Transaction(
|
|
wallet_id=wallet.id,
|
|
type=tx_type,
|
|
amount=amount,
|
|
balance_after=wallet.balance,
|
|
reference_type=reference_type,
|
|
reference_id=reference_id,
|
|
idempotency_key=idempotency_key,
|
|
metadata_json=metadata,
|
|
)
|
|
db.add(tx)
|
|
await db.commit()
|
|
await db.refresh(wallet)
|
|
await db.refresh(tx)
|
|
return wallet, tx
|
|
|
|
@classmethod
|
|
async def debit(
|
|
cls,
|
|
db: AsyncSession,
|
|
user_id: int,
|
|
amount: int,
|
|
tx_type: TransactionType | str,
|
|
reference_type: str | None = None,
|
|
reference_id: str | None = None,
|
|
idempotency_key: str | None = None,
|
|
metadata: dict[str, Any] | None = None,
|
|
) -> tuple[Wallet, Transaction]:
|
|
"""Debit a user's wallet atomically with row locking and balance check.
|
|
|
|
Raises InsufficientFundsError if wallet balance < amount.
|
|
If idempotency_key is provided and already exists, returns the original
|
|
outcome without double-debiting.
|
|
"""
|
|
if amount <= 0:
|
|
raise InvalidAmountError("Debit amount must be strictly positive")
|
|
|
|
if isinstance(tx_type, str):
|
|
tx_type = TransactionType(tx_type)
|
|
|
|
# Check idempotency first
|
|
if idempotency_key:
|
|
stmt = select(Transaction).where(Transaction.idempotency_key == idempotency_key)
|
|
existing_tx = (await db.execute(stmt)).scalar_one_or_none()
|
|
if existing_tx is not None:
|
|
wallet = await cls.get_or_create_wallet(db, user_id, for_update=False)
|
|
return wallet, existing_tx
|
|
|
|
# Lock wallet row
|
|
wallet = await cls.get_or_create_wallet(db, user_id, for_update=True)
|
|
cur_balance = wallet.balance
|
|
if cur_balance < amount:
|
|
await db.rollback()
|
|
raise InsufficientFundsError(
|
|
f"Insufficient funds: available {cur_balance}, required {amount}"
|
|
)
|
|
|
|
wallet.balance -= amount
|
|
|
|
tx = Transaction(
|
|
wallet_id=wallet.id,
|
|
type=tx_type,
|
|
amount=-amount,
|
|
balance_after=wallet.balance,
|
|
reference_type=reference_type,
|
|
reference_id=reference_id,
|
|
idempotency_key=idempotency_key,
|
|
metadata_json=metadata,
|
|
)
|
|
db.add(tx)
|
|
await db.commit()
|
|
await db.refresh(wallet)
|
|
await db.refresh(tx)
|
|
return wallet, tx
|
|
|
|
@staticmethod
|
|
async def get_transactions(
|
|
db: AsyncSession, wallet_id: int, limit: int = 50, offset: int = 0
|
|
) -> tuple[list[Transaction], int]:
|
|
"""Fetch paged transactions for a wallet ordered newest-first, plus total count."""
|
|
count_stmt = select(func.count()).select_from(Transaction).where(Transaction.wallet_id == wallet_id)
|
|
total = (await db.execute(count_stmt)).scalar() or 0
|
|
|
|
stmt = (
|
|
select(Transaction)
|
|
.where(Transaction.wallet_id == wallet_id)
|
|
.order_by(Transaction.created_at.desc(), Transaction.id.desc())
|
|
.offset(offset)
|
|
.limit(limit)
|
|
)
|
|
res = await db.execute(stmt)
|
|
return list(res.scalars().all()), total
|
|
|
|
|
|
wallet_service = WalletService()
|