38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
from __future__ import annotations
|
||
|
||
from collections.abc import Iterable
|
||
|
||
from sqlalchemy import text
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
|
||
async def _acquire_key_lock(db: AsyncSession, lock_key: str) -> None:
|
||
"""PostgreSQL事务级 advisory lock；SQLite 本地调试环境无需额外锁。"""
|
||
bind = db.get_bind()
|
||
dialect_name = bind.dialect.name if bind is not None else ""
|
||
if dialect_name == "postgresql":
|
||
await db.execute(
|
||
text("SELECT pg_advisory_xact_lock(hashtextextended(:lock_key, 0))"),
|
||
{"lock_key": lock_key},
|
||
)
|
||
|
||
|
||
async def acquire_user_credit_lock(db: AsyncSession, user_id: str) -> None:
|
||
await _acquire_key_lock(db, f"credit:user:{user_id}")
|
||
|
||
|
||
async def acquire_team_business_lock(db: AsyncSession, team_id: str) -> None:
|
||
await _acquire_key_lock(db, f"credit:team:{team_id}")
|
||
|
||
|
||
async def acquire_subscription_credit_lock(db: AsyncSession, subscription_id: str) -> None:
|
||
await _acquire_key_lock(db, f"credit:subscription:{subscription_id}")
|
||
|
||
|
||
async def acquire_subscription_credit_locks(
|
||
db: AsyncSession,
|
||
subscription_ids: Iterable[str],
|
||
) -> None:
|
||
for subscription_id in sorted({str(item) for item in subscription_ids if item}):
|
||
await acquire_subscription_credit_lock(db, subscription_id)
|