379 lines
10 KiB
Python
379 lines
10 KiB
Python
"""
|
||
余额仓库
|
||
|
||
处理余额相关的数据库操作。
|
||
|
||
设计说明：
|
||
- 使用乐观锁（version）防止并发更新冲突
|
||
- 提供行级锁支持（悲观锁）用于关键操作
|
||
"""
|
||
|
||
from datetime import datetime
|
||
from typing import Any
|
||
|
||
from sqlalchemy import select, update, and_
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.models.balance import (
|
||
UserBalance,
|
||
BalanceTransaction,
|
||
TransactionType,
|
||
TransactionStatus,
|
||
)
|
||
from app.repositories.base import BaseRepository
|
||
|
||
|
||
class BalanceRepository(BaseRepository[UserBalance]):
|
||
"""余额仓库"""
|
||
|
||
model = UserBalance
|
||
|
||
async def get_by_user_id(self, user_id: str) -> UserBalance | None:
|
||
"""
|
||
通过用户 ID 获取余额账户
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
|
||
Returns:
|
||
余额账户或 None
|
||
"""
|
||
stmt = select(UserBalance).where(UserBalance.user_id == user_id)
|
||
result = await self.session.execute(stmt)
|
||
return result.scalar_one_or_none()
|
||
|
||
async def get_by_user_id_for_update(self, user_id: str) -> UserBalance | None:
|
||
"""
|
||
通过用户 ID 获取余额账户（加行级锁）
|
||
|
||
用于需要原子性更新的场景，如扣款操作。
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
|
||
Returns:
|
||
余额账户或 None
|
||
"""
|
||
stmt = (
|
||
select(UserBalance)
|
||
.where(UserBalance.user_id == user_id)
|
||
.with_for_update() # 行级锁
|
||
)
|
||
result = await self.session.execute(stmt)
|
||
return result.scalar_one_or_none()
|
||
|
||
async def get_or_create(self, user_id: str) -> UserBalance:
|
||
"""
|
||
获取或创建余额账户
|
||
|
||
如果用户没有余额账户，自动创建一个。
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
|
||
Returns:
|
||
余额账户
|
||
"""
|
||
balance = await self.get_by_user_id(user_id)
|
||
if balance is None:
|
||
balance = await self.create(user_id=user_id)
|
||
return balance
|
||
|
||
async def get_or_create_for_update(self, user_id: str) -> UserBalance:
|
||
"""
|
||
获取或创建余额账户（加行级锁）
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
|
||
Returns:
|
||
余额账户
|
||
"""
|
||
balance = await self.get_by_user_id_for_update(user_id)
|
||
if balance is None:
|
||
balance = await self.create(user_id=user_id)
|
||
# 重新获取并加锁
|
||
await self.session.flush()
|
||
balance = await self.get_by_user_id_for_update(user_id)
|
||
return balance # type: ignore
|
||
|
||
async def update_balance_optimistic(
|
||
self,
|
||
balance: UserBalance,
|
||
delta: int,
|
||
*,
|
||
is_recharge: bool = False,
|
||
is_consumption: bool = False,
|
||
) -> bool:
|
||
"""
|
||
使用乐观锁更新余额
|
||
|
||
通过版本号检查确保并发安全。
|
||
|
||
Args:
|
||
balance: 余额账户
|
||
delta: 变化量（正数增加，负数减少）
|
||
is_recharge: 是否为充值
|
||
is_consumption: 是否为消费
|
||
|
||
Returns:
|
||
是否更新成功
|
||
"""
|
||
current_version = balance.version
|
||
new_balance = balance.balance + delta
|
||
|
||
# 构建更新语句
|
||
update_values: dict[str, Any] = {
|
||
"balance": new_balance,
|
||
"version": current_version + 1,
|
||
}
|
||
|
||
if is_recharge and delta > 0:
|
||
update_values["total_recharged"] = balance.total_recharged + delta
|
||
if is_consumption and delta < 0:
|
||
update_values["total_consumed"] = balance.total_consumed + abs(delta)
|
||
|
||
stmt = (
|
||
update(UserBalance)
|
||
.where(
|
||
and_(
|
||
UserBalance.id == balance.id,
|
||
UserBalance.version == current_version, # 乐观锁检查
|
||
)
|
||
)
|
||
.values(**update_values)
|
||
)
|
||
|
||
result = await self.session.execute(stmt)
|
||
|
||
if result.rowcount == 1:
|
||
# 更新成功，刷新对象
|
||
balance.balance = new_balance
|
||
balance.version = current_version + 1
|
||
if is_recharge and delta > 0:
|
||
balance.total_recharged += delta
|
||
if is_consumption and delta < 0:
|
||
balance.total_consumed += abs(delta)
|
||
return True
|
||
|
||
return False
|
||
|
||
async def freeze_balance(
|
||
self,
|
||
balance: UserBalance,
|
||
amount: int,
|
||
) -> bool:
|
||
"""
|
||
冻结余额
|
||
|
||
Args:
|
||
balance: 余额账户
|
||
amount: 冻结金额（正数）
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
if amount <= 0:
|
||
return False
|
||
if balance.available_balance < amount:
|
||
return False
|
||
|
||
current_version = balance.version
|
||
stmt = (
|
||
update(UserBalance)
|
||
.where(
|
||
and_(
|
||
UserBalance.id == balance.id,
|
||
UserBalance.version == current_version,
|
||
UserBalance.balance - UserBalance.frozen_balance >= amount,
|
||
)
|
||
)
|
||
.values(
|
||
frozen_balance=UserBalance.frozen_balance + amount,
|
||
version=current_version + 1,
|
||
)
|
||
)
|
||
|
||
result = await self.session.execute(stmt)
|
||
|
||
if result.rowcount == 1:
|
||
balance.frozen_balance += amount
|
||
balance.version = current_version + 1
|
||
return True
|
||
|
||
return False
|
||
|
||
async def unfreeze_balance(
|
||
self,
|
||
balance: UserBalance,
|
||
amount: int,
|
||
) -> bool:
|
||
"""
|
||
解冻余额
|
||
|
||
Args:
|
||
balance: 余额账户
|
||
amount: 解冻金额（正数）
|
||
|
||
Returns:
|
||
是否成功
|
||
"""
|
||
if amount <= 0:
|
||
return False
|
||
if balance.frozen_balance < amount:
|
||
return False
|
||
|
||
current_version = balance.version
|
||
stmt = (
|
||
update(UserBalance)
|
||
.where(
|
||
and_(
|
||
UserBalance.id == balance.id,
|
||
UserBalance.version == current_version,
|
||
UserBalance.frozen_balance >= amount,
|
||
)
|
||
)
|
||
.values(
|
||
frozen_balance=UserBalance.frozen_balance - amount,
|
||
version=current_version + 1,
|
||
)
|
||
)
|
||
|
||
result = await self.session.execute(stmt)
|
||
|
||
if result.rowcount == 1:
|
||
balance.frozen_balance -= amount
|
||
balance.version = current_version + 1
|
||
return True
|
||
|
||
return False
|
||
|
||
|
||
class TransactionRepository(BaseRepository[BalanceTransaction]):
|
||
"""交易记录仓库"""
|
||
|
||
model = BalanceTransaction
|
||
|
||
async def get_by_user_id(
|
||
self,
|
||
user_id: str,
|
||
*,
|
||
offset: int = 0,
|
||
limit: int = 20,
|
||
transaction_type: TransactionType | None = None,
|
||
status: TransactionStatus | None = None,
|
||
) -> list[BalanceTransaction]:
|
||
"""
|
||
获取用户的交易记录
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
offset: 偏移量
|
||
limit: 限制数量
|
||
transaction_type: 交易类型过滤
|
||
status: 状态过滤
|
||
|
||
Returns:
|
||
交易记录列表
|
||
"""
|
||
stmt = select(BalanceTransaction).where(
|
||
BalanceTransaction.user_id == user_id
|
||
)
|
||
|
||
if transaction_type:
|
||
stmt = stmt.where(BalanceTransaction.transaction_type == transaction_type)
|
||
if status:
|
||
stmt = stmt.where(BalanceTransaction.status == status)
|
||
|
||
stmt = (
|
||
stmt
|
||
.order_by(BalanceTransaction.created_at.desc())
|
||
.offset(offset)
|
||
.limit(limit)
|
||
)
|
||
|
||
result = await self.session.execute(stmt)
|
||
return list(result.scalars().all())
|
||
|
||
async def count_by_user_id(
|
||
self,
|
||
user_id: str,
|
||
*,
|
||
transaction_type: TransactionType | None = None,
|
||
status: TransactionStatus | None = None,
|
||
) -> int:
|
||
"""
|
||
统计用户的交易记录数量
|
||
|
||
Args:
|
||
user_id: 用户 ID
|
||
transaction_type: 交易类型过滤
|
||
status: 状态过滤
|
||
|
||
Returns:
|
||
记录数量
|
||
"""
|
||
from sqlalchemy import func
|
||
|
||
stmt = select(func.count()).select_from(BalanceTransaction).where(
|
||
BalanceTransaction.user_id == user_id
|
||
)
|
||
|
||
if transaction_type:
|
||
stmt = stmt.where(BalanceTransaction.transaction_type == transaction_type)
|
||
if status:
|
||
stmt = stmt.where(BalanceTransaction.status == status)
|
||
|
||
result = await self.session.execute(stmt)
|
||
return result.scalar() or 0
|
||
|
||
async def get_by_idempotency_key(
|
||
self,
|
||
idempotency_key: str,
|
||
) -> BalanceTransaction | None:
|
||
"""
|
||
通过幂等键获取交易记录
|
||
|
||
用于防止重复提交。
|
||
|
||
Args:
|
||
idempotency_key: 幂等键
|
||
|
||
Returns:
|
||
交易记录或 None
|
||
"""
|
||
stmt = select(BalanceTransaction).where(
|
||
BalanceTransaction.idempotency_key == idempotency_key
|
||
)
|
||
result = await self.session.execute(stmt)
|
||
return result.scalar_one_or_none()
|
||
|
||
async def get_by_reference(
|
||
self,
|
||
reference_type: str,
|
||
reference_id: str,
|
||
) -> list[BalanceTransaction]:
|
||
"""
|
||
通过业务关联获取交易记录
|
||
|
||
Args:
|
||
reference_type: 关联业务类型
|
||
reference_id: 关联业务 ID
|
||
|
||
Returns:
|
||
交易记录列表
|
||
"""
|
||
stmt = (
|
||
select(BalanceTransaction)
|
||
.where(
|
||
and_(
|
||
BalanceTransaction.reference_type == reference_type,
|
||
BalanceTransaction.reference_id == reference_id,
|
||
)
|
||
)
|
||
.order_by(BalanceTransaction.created_at.desc())
|
||
)
|
||
result = await self.session.execute(stmt)
|
||
return list(result.scalars().all())
|
||
|