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)