318 lines
13 KiB
Python
318 lines
13 KiB
Python
"""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 = "<unavailable>"
|
||
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()
|