sdpy-kit 0.1.0__py3-none-any.whl

kit/db/transaction.py

@@ -0,0 +1,109 @@
+"""
+Copyright (c) 2021-2026 Clark Chang. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+Project: sdpy
+Author: Clark Chang
+"""
+
+# date: 2026-08-23
+"""事务管理模块
+
+提供 @transactional 装饰器：捕获 SQLAlchemyError（含 DBError 及其所有子类），
+确保重试逻辑对数据库异常有效。
+"""
+
+import functools
+import inspect
+import time
+from collections.abc import Callable
+from typing import TypeVar
+
+from sqlalchemy.exc import SQLAlchemyError
+
+from kit.config.log_config import get_default_logger
+from kit.db.exceptions import DBError
+from kit.db.session import session_scope
+
+logger = get_default_logger(__name__)
+
+F = TypeVar("F", bound=Callable[..., object])
+
+
+def transactional(
+    db_type: str | None = None,
+    *,
+    retry_count: int = 0,
+    retry_delay: float = 0.5,
+) -> Callable[[F], F]:
+    """事务装饰器：自动注入 session，统一捕获 SQLAlchemyError 并重试
+
+    被装饰的函数必须接受 ``session`` 关键字参数。
+    捕获 SQLAlchemyError（含 DBError、IntegrityError、OperationalError 等所有子类）。
+
+    Args:
+        db_type: 数据库类型，None 时使用配置默认值
+        retry_count: 失败后的额外重试次数（0 表示不重试）
+        retry_delay: 初始重试延迟（秒），每次重试翻倍（指数退避）
+
+    用法::
+
+        @transactional()
+        def transfer(session: Session, from_id: int, to_id: int, amount: float):
+            # 多个 CRUD 操作在同一事务中
+            sender = user_crud.get_or_raise(session, from_id)
+            receiver = user_crud.get_or_raise(session, to_id)
+            sender.balance -= amount
+            receiver.balance += amount
+            # 退出时自动提交
+    """
+
+    def decorator(func: F) -> F:
+        sig = inspect.signature(func)
+
+        @functools.wraps(func)
+        def wrapper(*args: object, **kwargs: object) -> object:
+            # 检查函数是否接受 session 参数
+            if "session" not in sig.parameters:
+                raise TypeError(
+                    f"被 @transactional 装饰的函数 '{func.__name__}' "
+                    "必须接受 'session' 关键字参数"
+                )
+
+            last_exception: Exception | None = None
+
+            for attempt in range(retry_count + 1):
+                try:
+                    with session_scope(db_type, autocommit=True) as session:
+                        kwargs["session"] = session
+                        return func(*args, **kwargs)
+                except SQLAlchemyError as e:
+                    last_exception = e
+                    if attempt < retry_count:
+                        delay = retry_delay * (2**attempt)
+                        logger.warning(
+                            f"⚠️ 事务失败，{delay:.1f}s 后重试 "
+                            f"({attempt + 1}/{retry_count}): {e}"
+                        )
+                        time.sleep(delay)
+                    else:
+                        break
+
+            raise DBError(
+                f"事务执行失败（重试 {retry_count} 次后仍失败）: {last_exception}"
+            )
+
+        return wrapper  # type: ignore
+
+    return decorator