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}


@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
            if error_code not in MYSQL_RETRYABLE_ERROR_CODES or attempt == max_attempts:
                raise
            sleep(base_delay_seconds * attempt)

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