"""SQLAlchemy 事务适配器(共享基础设施)。 实现 ``yuxi.storage.transactions.ports.TransactionPort`` 契约,基于 SQLAlchemy ``AsyncSession`` 提供事务边界控制能力。本模块是 yuxi 三模块 (channels / external_systems / scheduler)事务控制的统一实现,吸收 channels ``SqlAlchemyTransactionAdapter`` 的最佳实践(autobegin 隐式事务 处理、``is_active`` 防御 ``ResourceClosedError``)与 external_systems ``SqlAlchemyUnitOfWork`` 的 savepoint 能力。 事务边界由应用层(管道或用例编排器)显式调用,被驱动适配器通过 ``TransactionContext`` 透传(C-I1)加入同一事务,**不得** 自主提交。 日志兼容性:``logger`` 为 ``Any``,可能为 structlog logger(提供 ``bind``) 或 stdlib ``logging.Logger``(fallback)。诊断上下文(``transaction_id`` / ``session_id`` / ``error`` / ``is_active``)格式化进消息字符串,不传 kwargs, 保证两种 logger 均可用;``transaction_id`` 额外通过 ``bind`` 注入 structlog 上下文(若 logger 提供 ``bind``)。 """ from __future__ import annotations import logging import time import uuid from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager from typing import Any from sqlalchemy.ext.asyncio import AsyncSession from yuxi.storage.transactions.exceptions import ( TransactionBeginError, TransactionCommitError, TransactionRollbackError, ) from yuxi.storage.transactions.ports import TransactionPort __all__ = [ "SqlAlchemyTransactionAdapter", "SqlAlchemyTransactionContext", "with_session", ] _LOGGER = logging.getLogger(__name__) class SqlAlchemyTransactionContext: """SQLAlchemy 事务上下文。 封装 ``AsyncSession`` 与其事务。被驱动适配器在构造时共享同一 ``AsyncSession``,``begin()`` 在共享 session 上开启事务,所有适配器 的写操作自动加入。上下文管理器退出时自动提交(无异常)或回滚 (有异常);``commit`` / ``rollback`` 幂等。 诊断:构造时生成 ``transaction_id``,关键节点(``__aenter__`` / ``commit`` / ``rollback`` / ``savepoint``)记录 DEBUG 日志,失败时 记录 ERROR 日志,消息携带 ``transaction_id`` / ``session_id`` / ``is_active``。 被驱动适配器 **不得** 调用 ``commit`` / ``rollback``,仅由应用层 通过 ``TransactionPort`` 控制。 """ def __init__(self, session: AsyncSession, logger: Any | None = None) -> None: """初始化事务上下文。 Args: session: SQLAlchemy 异步会话,事务边界由本上下文控制。 logger: 可选日志器。提供 ``bind(transaction_id=...)`` 方法时 (如 structlog logger)注入 transaction_id 上下文;否则 直接使用。为 ``None`` 时回退到 ``logging.getLogger(__name__)``。 """ self._session = session self._txn: Any = None self.transaction_id = f"tx_{time.time_ns()}_{uuid.uuid4().hex[:8]}" if logger is None: logger = _LOGGER if hasattr(logger, "bind"): logger = logger.bind(transaction_id=self.transaction_id) self._logger = logger def get_session(self) -> AsyncSession: """返回底层共享 ``AsyncSession``,供未在构造时共享 session 的适配器复用主事务。 C-I1:``AgentRunAdapter`` 等适配器在构造时未与事务适配器共享 session, 通过本方法获取主事务的共享 session,复用同一事务,避免独立提交产生 孤儿记录。 """ return self._session async def commit(self) -> None: """提交当前事务(幂等)。 ``self._txn`` 为空时 no-op;非空时提交并清空引用。提交失败抛 ``TransactionCommitError`` 并通过 ``raise ... from exc`` 保留原始 异常链;此时事务引用不清空,由 ``__aexit__`` 按 ``is_active`` 判断后续处理。 """ if self._txn is None: return try: await self._txn.commit() except Exception as exc: self._logger.error( f"transaction commit failed | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={exc}" ) raise TransactionCommitError( "事务提交失败", transaction_id=self.transaction_id, session_state=self._session_state(), ) from exc is_active = self._txn.is_active self._txn = None self._logger.debug( f"transaction committed | transaction_id={self.transaction_id} " f"session_id={id(self._session)} is_active={is_active}" ) async def rollback(self) -> None: """回滚当前事务(幂等)。 ``self._txn`` 为空时 no-op;非空时回滚并清空引用。回滚失败抛 ``TransactionRollbackError`` 并保留原始异常链。 """ if self._txn is None: return try: await self._txn.rollback() except Exception as exc: self._logger.error( f"transaction rollback failed | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={exc}" ) raise TransactionRollbackError( "事务回滚失败", transaction_id=self.transaction_id, session_state=self._session_state(), ) from exc is_active = self._txn.is_active self._txn = None self._logger.debug( f"transaction rolled back | transaction_id={self.transaction_id} " f"session_id={id(self._session)} is_active={is_active}" ) @asynccontextmanager async def savepoint(self) -> AsyncIterator[None]: """创建 SAVEPOINT,异常时回滚到 SAVEPOINT 而不污染外层事务。 基于 ``self._session.begin_nested()`` 实现。正常退出时释放 SAVEPOINT,异常时由 ``begin_nested`` 上下文管理器回滚到 SAVEPOINT 并向上抛出原异常。用于批量操作中隔离每次迭代,避免单条失败导致 整个会话中毒(``PendingRollbackError``)。 """ self._logger.debug( f"savepoint begin | transaction_id={self.transaction_id} " f"session_id={id(self._session)} is_active={self._txn.is_active if self._txn is not None else None}" ) try: async with self._session.begin_nested(): yield except Exception as exc: self._logger.debug( f"savepoint rolled back | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={exc}" ) raise async def __aenter__(self) -> SqlAlchemyTransactionContext: """进入事务上下文,开启 SQLAlchemy 事务。 SQLAlchemy 2.0 autobegin 语义下,前置读操作(如 SELECT 查询)会在 session 上隐式开启只读事务。若不处理,``begin()`` 会抛 ``InvalidRequestError: A transaction is already begun on this Session.``。 此处检测并提交隐式事务后再开启显式事务。安全性保证:按适配器契约, ``tx=None`` 的写操作自主提交,autobegin 事务仅含读操作,提交不会 产生副作用数据落库。应用层显式事务边界(§10.1)由此方法独占控制。 失败时抛 ``TransactionBeginError`` 并保留原始异常链。 """ try: if self._session.in_transaction(): await self._session.commit() self._txn = await self._session.begin() except Exception as exc: self._logger.error( f"transaction begin failed | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={exc}" ) raise TransactionBeginError( "事务开启失败", transaction_id=self.transaction_id, session_state=self._session_state(), ) from exc self._logger.debug( f"transaction begun | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} is_active={self._txn.is_active}" ) return self async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None: """退出事务上下文。 若事务未被显式提交/回滚(``self._txn`` 非空且 ``is_active``),则按 异常状态自动提交(无异常)或回滚(有异常);若已被显式提交/回滚 (``self._txn`` 为空)或已被 SQLAlchemy 内部关闭(``is_active=False``, 如 flush 失败自动关闭事务),则跳过,避免 ``ResourceClosedError``。 事务对象退出后释放引用,避免泄漏。提交/回滚失败分别抛 ``TransactionCommitError`` / ``TransactionRollbackError`` 并保留 原始异常链。 """ try: if self._txn is not None and self._txn.is_active: if exc is None: try: await self._txn.commit() except Exception as e: self._logger.error( f"transaction commit failed on exit | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={e}" ) raise TransactionCommitError( "事务提交失败", transaction_id=self.transaction_id, session_state=self._session_state(), ) from e else: try: await self._txn.rollback() except Exception as e: self._logger.error( f"transaction rollback failed on exit | " f"transaction_id={self.transaction_id} " f"session_id={id(self._session)} error={e}" ) raise TransactionRollbackError( "事务回滚失败", transaction_id=self.transaction_id, session_state=self._session_state(), ) from e finally: self._txn = None def _session_state(self) -> str: """捕获当前 session/txn 状态用于异常诊断。""" txn_active = self._txn.is_active if self._txn is not None else None try: in_txn = self._session.in_transaction() except Exception: in_txn = "" return f"in_transaction={in_txn}, txn_active={txn_active}" class SqlAlchemyTransactionAdapter(TransactionPort): """SQLAlchemy 事务适配器。 实现 ``TransactionPort`` 契约,基于共享的 ``AsyncSession`` 提供事务 边界控制。事务边界由应用层显式调用 ``begin()`` 开启,被驱动适配器 通过构造时共享的 session 加入同一事务。 关键约束: - 事务边界 **必须** 由应用层控制。 - 被驱动适配器 **不得** 自主调用 ``commit`` / ``rollback``。 - 事务范围 **必须** 由应用层显式声明。 """ def __init__(self, session: AsyncSession, logger: Any | None = None) -> None: """初始化事务适配器。 Args: session: SQLAlchemy 异步会话,与被驱动适配器共享以保证事务 一致性。 logger: 可选日志器,透传给 ``SqlAlchemyTransactionContext`` 用于诊断日志。为 ``None`` 时回退到 ``logging.getLogger(__name__)``。 """ self._session = session self._logger = logger if logger is not None else _LOGGER def begin(self) -> SqlAlchemyTransactionContext: """开启一个新事务,返回 SQLAlchemy 事务上下文。 Returns: SQLAlchemy 事务上下文,被驱动适配器通过构造时共享的 session 加入同一事务。上下文管理器退出时自动提交(无异常)或回滚 (有异常)。 """ return SqlAlchemyTransactionContext(self._session, self._logger) @asynccontextmanager async def with_session( session_factory: Callable[[], AsyncSession], ) -> AsyncIterator[AsyncSession]: """应用级单例组件临时 session 获取辅助方法。 用完即关,不持有进程级 ``AsyncSession``。供未在构造时共享请求级 session 的应用级单例组件(如调度器、后台 worker)按需获取一次性 session。 Args: session_factory: session 工厂,每次调用返回新的 ``AsyncSession``。 Yields: 临时 ``AsyncSession``,退出时自动关闭。 """ session = session_factory() try: yield session finally: await session.close()