document-management/backend/dms/database/transaction.py

"""DMS事务、行锁和并发扩展入口。"""

from __future__ import annotations

from collections.abc import Callable, Generator
from contextlib import contextmanager
from time import sleep
from typing import Any, TypeVar

from sqlalchemy import Select
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import Session

from dms.extensions import db

T = TypeVar("T")
MYSQL_RETRYABLE_ERROR_CODES = {1205, 1213}
# 达梦等非MySQL驱动错误码不统一，死锁/锁超时以异常消息文本兜底识别。
DEADLOCK_MESSAGE_MARKERS = ("死锁", "deadlock", "DeadLock", "lock wait timeout")


@contextmanager
def transaction() -> Generator[Session, None, None]:
    """提供统一提交和异常回滚边界。"""
    session = db.session
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise


def for_update(statement: Select[Any]) -> Select[Any]:
    """为后续业务服务提供 ``SELECT ... FOR UPDATE`` 入口。"""
    return statement.with_for_update()


def execute_with_deadlock_retry(
    operation: Callable[[], T],
    *,
    max_attempts: int = 3,
    base_delay_seconds: float = 0.05,
) -> T:
    """重试数据库死锁与锁等待超时（MySQL错误码或消息特征），供业务事务显式调用。"""
    if max_attempts < 1:
        raise ValueError("max_attempts必须大于等于1")

    for attempt in range(1, max_attempts + 1):
        try:
            return operation()
        except OperationalError as exc:
            db.session.rollback()
            error_code = exc.orig.args[0] if getattr(exc.orig, "args", ()) else None
            retryable = error_code in MYSQL_RETRYABLE_ERROR_CODES or any(
                marker in str(exc.orig) for marker in DEADLOCK_MESSAGE_MARKERS
            )
            if not retryable or attempt == max_attempts:
                raise
            sleep(base_delay_seconds * attempt)

    raise RuntimeError("不可达的事务重试状态")