85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from fastapi import HTTPException
|
|
from sqlalchemy import text
|
|
from sqlalchemy.exc import DBAPIError, OperationalError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.config import settings
|
|
|
|
|
|
class DatabaseRowLockBusy(HTTPException, RuntimeError):
|
|
"""A short-lived PostgreSQL row/table lock could not be acquired in time."""
|
|
|
|
def __init__(self, message: str = "当前任务正在被其他流程处理，请稍后重试") -> None:
|
|
super().__init__(status_code=409, detail=message)
|
|
|
|
|
|
_LOCK_NOT_AVAILABLE_SQLSTATE = "55P03"
|
|
|
|
|
|
def _sqlstate_from_exception(exc: BaseException | None) -> str | None:
|
|
current: BaseException | None = exc
|
|
seen: set[int] = set()
|
|
while current is not None and id(current) not in seen:
|
|
seen.add(id(current))
|
|
for attr in ("sqlstate", "pgcode"):
|
|
value = getattr(current, attr, None)
|
|
if value:
|
|
return str(value)
|
|
current = getattr(current, "orig", None) or getattr(current, "__cause__", None)
|
|
return None
|
|
|
|
|
|
def is_postgres_lock_timeout(exc: BaseException) -> bool:
|
|
if _sqlstate_from_exception(exc) == _LOCK_NOT_AVAILABLE_SQLSTATE:
|
|
return True
|
|
message = str(exc).lower()
|
|
return "lock timeout" in message or "could not obtain lock" in message
|
|
|
|
|
|
def raise_if_database_lock_busy(exc: BaseException) -> None:
|
|
if is_postgres_lock_timeout(exc):
|
|
raise DatabaseRowLockBusy("数据库任务行正在被其他事务处理，请稍后重试") from exc
|
|
|
|
|
|
async def apply_short_lock_timeout(
|
|
db: AsyncSession,
|
|
*,
|
|
seconds: int | None = None,
|
|
) -> None:
|
|
"""Apply a transaction-local PostgreSQL lock wait limit.
|
|
|
|
It deliberately does not change the global database configuration and is a
|
|
no-op on non-PostgreSQL test/development databases.
|
|
"""
|
|
bind = db.get_bind()
|
|
if bind is None or bind.dialect.name != "postgresql":
|
|
return
|
|
timeout_seconds = max(
|
|
1,
|
|
int(
|
|
seconds
|
|
if seconds is not None
|
|
else getattr(settings, "GENERATION_DB_LOCK_TIMEOUT_SECONDS", 5)
|
|
or 5
|
|
),
|
|
)
|
|
await db.execute(text(f"SET LOCAL lock_timeout = '{timeout_seconds}s'"))
|
|
|
|
|
|
async def execute_with_lock_timeout(
|
|
db: AsyncSession,
|
|
statement: Any,
|
|
*,
|
|
seconds: int | None = None,
|
|
):
|
|
await apply_short_lock_timeout(db, seconds=seconds)
|
|
try:
|
|
return await db.execute(statement)
|
|
except (OperationalError, DBAPIError) as exc:
|
|
raise_if_database_lock_busy(exc)
|
|
raise
|