Implement WalletService with atomic credit/debit, row locking, and wallet/transactions API (task 909e3e9a59af)
This commit is contained in:
@@ -0,0 +1,199 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user