以下方案面向“数据库连接池-事务上下文”的资源管理需求，强调高可靠的资源获取与归还、事务自动提交/回滚、语句级超时、可配置隔离级别以及可观测性。代码遵循Python最佳实践、具备完善异常处理，并提供重试策略的配套执行器（对事务体进行重试需要外部函数式封装，而不是仅依赖with语法）。
- 技术需求分析
-
资源类型与关键挑战
- 连接池获取/归还：
- 必须在任何路径（正常/异常/超时/中断）都将连接安全归还到池中，避免泄漏。
- 连接获取需要支持超时，若底层池不支持原生超时，需要在适配层进行轮询+睡眠直到超时。
- 事务管理：
- 进入上下文自动开启事务（BEGIN），设置隔离级别（默认 REPEATABLE READ），可配置只读/读写。
- 成功路径自动提交；异常路径自动回滚，且不掩盖原始异常。
- commit 失败时需要尝试回滚连接状态并向上抛出 commit 异常（不丢失信息）。
- 语句级超时：
- 在事务内使用 SET LOCAL statement_timeout 实现语句级超时（以毫秒为单位，PostgreSQL）。
- 同理可选设置 lock_timeout、idle_in_transaction_session_timeout。
- 重试策略：
- 事务体的重试不能在单个上下文管理器内部完成（with 语句块无法自动重放），需要提供一个“执行器”函数：按策略在每次尝试内创建一个新的事务上下文并执行用户函数。
- 需要可插拔的“可重试异常判定器”，例如 PostgreSQL 的 SerializationFailure、DeadlockDetected 或 SQLSTATE 类别40P01/40001。
- 可观测性：
- 记录事务开始/提交/回滚/重试的日志、耗时、异常的 SQLSTATE（若可用）。
- 留出度量与追踪Hook。
-
线程安全与性能
- 每个事务上下文实例只在当前线程使用；不共享可变状态。
- 尽量使用 SET LOCAL 限定作用域到当前事务，避免跨事务影响。
- 游标资源由会话对象跟踪并在退出时关闭，防止游标泄漏。
- 上下文管理器类实现
下面代码包含：
- PoolAdapter 协议与两个常用适配器（队列池、psycopg2 pool）
- RetryPolicy 重试策略与回退算法
- TransactionSession：包装连接并跟踪游标
- TransactionContext：上下文管理器（单次事务）
- run_in_transaction：基于重试策略的事务执行器
- PostgreSQL 的默认可重试异常判定器
注：示例默认面向 PostgreSQL（psycopg2/psycopg），但接口层尽量DB-API 2.0泛化；对于其他数据库可替换判定器与设置语句。
from __future__ import annotations
import logging
import time
import random
from dataclasses import dataclass
from enum import Enum
from typing import (
Any,
Callable,
ContextManager,
Iterable,
Optional,
Protocol,
Sequence,
Tuple,
Type,
TypeVar,
Union,
runtime_checkable,
Generic,
List,
Set,
)
logger = logging.getLogger(__name__)
T = TypeVar("T")
# ---------- DB-API minimal protocols ----------
@runtime_checkable
class DBCursor(Protocol):
def execute(self, operation: str, params: Optional[Union[Sequence[Any], dict[str, Any]]] = None) -> Any: ...
def close(self) -> None: ...
@property
def rowcount(self) -> int: ...
def fetchone(self) -> Any: ...
def fetchall(self) -> list[Any]: ...
@runtime_checkable
class DBConnection(Protocol):
def cursor(self) -> DBCursor: ...
def commit(self) -> None: ...
def rollback(self) -> None: ...
# autocommit 为可选属性，部分驱动支持（如 psycopg2/3）
autocommit: Any # type: ignore[attr-defined]
# ---------- Pool adapter protocol & implementations ----------
class PoolError(Exception):
"""Raised when the pool cannot provide a connection within the configured constraints."""
@runtime_checkable
class PoolAdapter(Protocol):
"""
Adapter over a connection pool, providing acquire/release with optional timeout.
Implementations must not leak connections: after acquire success, release(conn) must be callable.
"""
def acquire(self, timeout: Optional[float] = None) -> DBConnection: ...
def release(self, conn: DBConnection) -> None: ...
class QueuePoolAdapter(PoolAdapter):
"""
Wraps a queue.Queue-like pool (get(timeout=...), put(conn)) into PoolAdapter.
"""
def __init__(self, queue_obj: Any):
self._q = queue_obj
def acquire(self, timeout: Optional[float] = None) -> DBConnection:
import queue
try:
return self._q.get(timeout=timeout) # type: ignore[no-any-return]
except queue.Empty as e:
raise PoolError("Acquire connection timeout from queue pool") from e
def release(self, conn: DBConnection) -> None:
self._q.put(conn)
class Psycopg2PoolAdapter(PoolAdapter):
"""
Wraps psycopg2.pool.ThreadedConnectionPool-like object.
Behavior:
- psycopg2 pool doesn't offer native timeout; this adapter polls until deadline.
- When pool is exhausted, getconn() may raise, we retry until timeout.
"""
def __init__(self, psycopg2_pool: Any, poll_interval: float = 0.05):
self._pool = psycopg2_pool
self._poll_interval = poll_interval
def acquire(self, timeout: Optional[float] = None) -> DBConnection:
deadline = time.monotonic() + timeout if timeout and timeout > 0 else None
last_exc: Optional[BaseException] = None
while True:
try:
conn = self._pool.getconn()
if conn is None:
raise PoolError("psycopg2 pool returned None")
return conn
except Exception as e: # Pool exhausted or other issues
last_exc = e
if deadline is not None and time.monotonic() >= deadline:
raise PoolError("Acquire connection timeout from psycopg2 pool") from last_exc
time.sleep(self._poll_interval)
def release(self, conn: DBConnection) -> None:
self._pool.putconn(conn)
# ---------- Isolation levels ----------
class IsolationLevel(str, Enum):
READ_COMMITTED = "READ COMMITTED"
REPEATABLE_READ = "REPEATABLE READ"
SERIALIZABLE = "SERIALIZABLE"
# READ UNCOMMITTED 对PostgreSQL等同 READ COMMITTED；如需支持可自行映射
# READ_UNCOMMITTED = "READ UNCOMMITTED"
# ---------- Retry policy ----------
@dataclass(frozen=True)
class RetryPolicy:
max_attempts: int = 3
initial_backoff: float = 0.05 # seconds
max_backoff: float = 0.8 # seconds
jitter: float = 0.1 # +/- randomization
deadline: Optional[float] = None # absolute seconds since epoch or None
retriable_predicate: Callable[[BaseException], bool] = lambda e: False
def compute_backoff(self, attempt: int) -> float:
# Exponential backoff with jitter
base = min(self.max_backoff, self.initial_backoff * (2 ** (attempt - 1)))
if self.jitter > 0:
delta = random.uniform(-self.jitter, self.jitter) * base
return max(0.0, base + delta)
return base
# ---------- Session wrapper to track cursors ----------
class TransactionSession:
"""
A helper that wraps a DB-API connection within a transaction.
Tracks cursors created via this object and closes them on context exit.
"""
def __init__(self, conn: DBConnection, name: Optional[str] = None):
self._conn = conn
self._cursors: Set[DBCursor] = set()
self.name = name
@property
def connection(self) -> DBConnection:
return self._conn
def cursor(self) -> DBCursor:
cur = self._conn.cursor()
self._cursors.add(cur)
return cur
def execute(self, sql: str, params: Optional[Union[Sequence[Any], dict[str, Any]]] = None) -> DBCursor:
cur = self.cursor()
cur.execute(sql, params)
return cur
def _close_all_cursors(self) -> None:
for cur in list(self._cursors):
try:
cur.close()
except Exception: # log but don't raise
logger.warning("Failed to close cursor for session %s", self.name, exc_info=True)
finally:
self._cursors.discard(cur)
# ---------- Transaction context manager ----------
class TransactionContext(ContextManager[TransactionSession]):
"""
Context manager for a single database transaction from a pool.
Features:
- Acquire connection with optional timeout and release it safely.
- BEGIN transaction with configurable isolation level and read_only flag.
- SET LOCAL statement_timeout / lock_timeout if provided.
- Commit on success; rollback on failure, without masking original exception.
- Tracks cursors created via session and closes them at exit.
- Restores original autocommit if supported by the driver.
Note: This is a single-attempt transaction manager.
For retries, use run_in_transaction() which re-invokes the body per attempt.
"""
def __init__(
self,
pool: PoolAdapter,
*,
acquire_timeout: Optional[float] = None,
isolation_level: IsolationLevel = IsolationLevel.REPEATABLE_READ,
read_only: bool = False,
statement_timeout_ms: Optional[int] = None,
lock_timeout_ms: Optional[int] = None,
idle_in_txn_session_timeout_ms: Optional[int] = None,
name: Optional[str] = None,
on_event: Optional[Callable[[str, dict[str, Any]], None]] = None,
) -> None:
self._pool = pool
self._acquire_timeout = acquire_timeout
self._isolation_level = isolation_level
self._read_only = read_only
self._statement_timeout_ms = statement_timeout_ms
self._lock_timeout_ms = lock_timeout_ms
self._idle_in_txn_session_timeout_ms = idle_in_txn_session_timeout_ms
self._name = name or "txn"
self._on_event = on_event
self._conn: Optional[DBConnection] = None
self._sess: Optional[TransactionSession] = None
self._old_autocommit: Optional[bool] = None
self._start_ts: Optional[float] = None
def __enter__(self) -> TransactionSession:
self._start_ts = time.monotonic()
self._emit("acquire.start", {"name": self._name})
conn = self._pool.acquire(timeout=self._acquire_timeout)
self._conn = conn
self._emit("acquire.ok", {"name": self._name})
# autocommit handling (best-effort; may not exist in some drivers)
self._old_autocommit = getattr(conn, "autocommit", None)
try:
if self._old_autocommit is not None:
try:
# Ensure explicit transaction boundaries
conn.autocommit = False # type: ignore[assignment]
except Exception:
# If driver doesn't support, ignore but log
logger.debug("Driver does not support autocommit property change.", exc_info=True)
# Begin transaction explicitly for clarity
cur = conn.cursor()
try:
# Set transaction characteristics
# Isolation level
cur.execute(f"BEGIN")
cur.execute(f"SET TRANSACTION ISOLATION LEVEL {self._isolation_level.value}")
# Read-only / read-write
cur.execute("SET TRANSACTION READ ONLY" if self._read_only else "SET TRANSACTION READ WRITE")
# Transaction-local timeouts (PostgreSQL)
if self._statement_timeout_ms is not None:
cur.execute("SET LOCAL statement_timeout = %s", (self._statement_timeout_ms,))
if self._lock_timeout_ms is not None:
cur.execute("SET LOCAL lock_timeout = %s", (self._lock_timeout_ms,))
if self._idle_in_txn_session_timeout_ms is not None:
cur.execute("SET LOCAL idle_in_transaction_session_timeout = %s",
(self._idle_in_txn_session_timeout_ms,))
finally:
try:
cur.close()
except Exception:
logger.warning("Failed to close internal cursor when starting transaction.", exc_info=True)
self._sess = TransactionSession(conn, name=self._name)
self._emit("begin.ok", {"name": self._name, "isolation": self._isolation_level.value,
"read_only": self._read_only})
return self._sess
except Exception:
# If any failure before returning session, ensure connection is returned
self._safe_release()
self._emit("begin.error", {"name": self._name}, level="error")
raise
def __exit__(self, exc_type, exc, tb) -> bool:
assert self._conn is not None, "Invariant: conn must exist on __exit__"
# Always close user cursors created via session
if self._sess is not None:
self._sess._close_all_cursors()
commit_ok = False
try:
if exc_type is None:
# commit path
self._emit("commit.start", {"name": self._name})
try:
self._conn.commit()
commit_ok = True
self._emit("commit.ok", {"name": self._name})
except Exception as commit_exc:
# Try rollback to clean connection state, but do not mask commit exception
self._emit("commit.error", {"name": self._name, "error": repr(commit_exc)}, level="error")
try:
self._conn.rollback()
except Exception:
logger.error("Rollback after commit failure also failed.", exc_info=True)
raise
else:
# rollback path
self._emit("rollback.start", {"name": self._name, "exc": repr(exc)})
try:
self._conn.rollback()
self._emit("rollback.ok", {"name": self._name})
except Exception:
# Must not mask the original exception from the with-body
self._emit("rollback.error", {"name": self._name}, level="error")
logger.error("Rollback failed; original exception will be re-raised.", exc_info=True)
# return False -> re-raise original exception
return False
finally:
# Restore autocommit if supported
try:
if self._old_autocommit is not None:
self._conn.autocommit = self._old_autocommit # type: ignore[assignment]
except Exception:
logger.debug("Failed to restore autocommit.", exc_info=True)
# Release connection back to pool in any case
self._safe_release()
duration = (time.monotonic() - self._start_ts) if self._start_ts else None
self._emit("finish", {"name": self._name, "duration_s": duration, "commit_ok": commit_ok})
# On commit path, returning False will propagate any commit exception (already raised), else True.
return True
def _safe_release(self) -> None:
if self._conn is not None:
try:
self._pool.release(self._conn)
except Exception:
# Pool release should never raise to callers; just log.
logger.error("Failed to release connection to pool.", exc_info=True)
finally:
self._conn = None
def _emit(self, event: str, payload: dict[str, Any], level: str = "debug") -> None:
payload = dict(payload)
payload["event"] = event
if self._on_event:
try:
self._on_event(event, payload)
except Exception:
logger.debug("on_event hook failed.", exc_info=True)
# Fallback logging
if level == "error":
logger.error("%s | %s", event, payload)
elif level == "info":
logger.info("%s | %s", event, payload)
else:
logger.debug("%s | %s", event, payload)
# ---------- Retriable predicate helpers (PostgreSQL) ----------
def is_postgres_retriable(exc: BaseException) -> bool:
"""
Try to detect serialization/deadlock errors in a driver-agnostic way for PostgreSQL.
- psycopg2/3 exceptions may expose pgcode / sqlstate or specific classes.
- We inspect common attributes and class names as a last resort.
"""
# Check SQLSTATE if available
sqlstate = getattr(exc, "pgcode", None) or getattr(exc, "sqlstate", None)
if isinstance(sqlstate, str) and sqlstate in ("40001", "40P01"): # serialization_failure, deadlock_detected
return True
# psycopg3 exceptions may live in psycopg.errors with specific types
cls_name = exc.__class__.__name__
if cls_name in ("SerializationFailure", "DeadlockDetected"):
return True
# Fall back: heuristic on message (avoid relying on message ideally)
msg = str(exc).lower()
if "deadlock detected" in msg or "could not serialize access" in msg:
return True
return False
# ---------- Transaction executor with retry ----------
def run_in_transaction(
pool: PoolAdapter,
fn: Callable[[TransactionSession], T],
*,
isolation_level: IsolationLevel = IsolationLevel.REPEATABLE_READ,
read_only: bool = False,
statement_timeout_ms: Optional[int] = None,
lock_timeout_ms: Optional[int] = None,
idle_in_txn_session_timeout_ms: Optional[int] = None,
acquire_timeout: Optional[float] = None,
retry: RetryPolicy = RetryPolicy(max_attempts=3, retriable_predicate=is_postgres_retriable),
name: Optional[str] = None,
on_event: Optional[Callable[[str, dict[str, Any]], None]] = None,
) -> T:
"""
Execute fn(session) within a transaction, retrying on retriable errors as per policy.
Returns:
The return value of fn from the successful attempt.
Raises:
The last exception if all retries fail or non-retriable error occurs.
"""
attempt = 0
last_exc: Optional[BaseException] = None
start_wall = time.time()
while True:
attempt += 1
if retry.deadline is not None and time.time() > retry.deadline:
# Deadline exceeded before starting attempt
if last_exc:
raise last_exc
raise TimeoutError("Retry deadline exceeded before starting transaction attempt")
attempt_name = f"{name or 'txn'}#attempt{attempt}"
try:
with TransactionContext(
pool,
acquire_timeout=acquire_timeout,
isolation_level=isolation_level,
read_only=read_only,
statement_timeout_ms=statement_timeout_ms,
lock_timeout_ms=lock_timeout_ms,
idle_in_txn_session_timeout_ms=idle_in_txn_session_timeout_ms,
name=attempt_name,
on_event=on_event,
) as sess:
return fn(sess)
except BaseException as e:
last_exc = e
# Non-retriable or attempts exhausted?
should_retry = retry.retriable_predicate(e)
if not should_retry or attempt >= retry.max_attempts:
# Include timing info in logs
elapsed = time.time() - start_wall
logger.error("Transaction failed after %d attempt(s), elapsed=%.3fs", attempt, elapsed, exc_info=True)
raise
# Backoff and continue
backoff = retry.compute_backoff(attempt)
if retry.deadline is not None and time.time() + backoff > retry.deadline:
logger.warning("Backoff would exceed retry deadline; aborting retries.", exc_info=True)
raise
logger.info("Retriable transaction error on attempt %d; backing off %.3fs: %r", attempt, backoff, e)
time.sleep(backoff)
continue
- 使用示例
示例1：基础用法（单次事务，不带重试）
import psycopg2
from psycopg2.pool import ThreadedConnectionPool
pool = ThreadedConnectionPool(minconn=5, maxconn=20, dsn="postgresql://user:pwd@host:5432/db")
adapter = Psycopg2PoolAdapter(pool)
def create_user(sess: TransactionSession, uid: int, name: str) -> None:
sess.execute("INSERT INTO users(id, name) VALUES (%s, %s)", (uid, name))
with TransactionContext(
adapter,
isolation_level=IsolationLevel.REPEATABLE_READ,
statement_timeout_ms=1500,
acquire_timeout=2.0,
name="create_user_txn",
) as session:
create_user(session, 1001, "Alice")
示例2：带重试策略的订单创建（多表写入 + outbox，确保一致性）
- 事务内只写数据库（订单表、订单项表、消息 outbox 表），提交后由异步消费者读取 outbox 并发布到消息队列，确保最终一致性。
from typing import Any
def create_order_txn(sess: TransactionSession, payload: dict[str, Any]) -> int:
# 写入订单主表
cur = sess.execute(
"INSERT INTO orders(user_id, total_amount, status) VALUES (%s, %s, %s) RETURNING id",
(payload["user_id"], payload["total_amount"], "CREATED")
)
order_id = cur.fetchone()[0]
# 写入订单项
for item in payload["items"]:
sess.execute(
"INSERT INTO order_items(order_id, sku, qty, price) VALUES (%s, %s, %s, %s)",
(order_id, item["sku"], item["qty"], item["price"])
)
# 写入 outbox 消息（事务内）
sess.execute(
"INSERT INTO outbox(aggregate, aggregate_id, event_type, payload) VALUES (%s, %s, %s, %s)",
("order", order_id, "OrderCreated", payload)
)
return order_id
retry_policy = RetryPolicy(
max_attempts=5,
initial_backoff=0.05,
max_backoff=0.5,
jitter=0.2,
retriable_predicate=is_postgres_retriable,
)
order_id = run_in_transaction(
adapter,
lambda s: create_order_txn(s, {
"user_id": 42,
"total_amount": 399.0,
"items": [{"sku": "A100", "qty": 1, "price": 199.0}, {"sku": "B200", "qty": 2, "price": 100.0}],
}),
isolation_level=IsolationLevel.REPEATABLE_READ,
statement_timeout_ms=2000,
acquire_timeout=2.0,
retry=retry_policy,
name="create_order",
)
print("order_id=", order_id)
# 提交成功后，异步进程（或后台任务）扫描 outbox 表并发布到消息队列，保证最终一致性与可观测性。
示例3：只读查询 + 每语句超时
def read_report(sess: TransactionSession) -> list[tuple]:
cur = sess.execute("SELECT id, total_amount FROM orders WHERE status = %s", ("CREATED",))
return cur.fetchall()
result = run_in_transaction(
adapter,
read_report,
isolation_level=IsolationLevel.REPEATABLE_READ,
read_only=True,
statement_timeout_ms=500, # 语句级超时
acquire_timeout=1.0,
retry=RetryPolicy(max_attempts=2, retriable_predicate=is_postgres_retriable),
name="report_ro",
)
- 技术原理说明
- 资源获取与释放
- 通过 PoolAdapter 统一连接池接口，适配多个来源（队列池、psycopg2 pool等）。
- 上下文进入时 acquire，退出时无论成功/失败均 release，严格防止连接泄漏。
- 事务控制
- 显示 BEGIN + SET TRANSACTION 控制隔离级别与读写模式，可重复读默认提高读一致性。
- 使用 SET LOCAL 设置 statement_timeout / lock_timeout 等，仅作用于当前事务生命周期，防止污染连接会话。
- 成功提交路径：commit；异常路径：rollback 且不掩盖原始异常（exit 返回 False）。
- commit 失败时先尝试 rollback 以清理连接状态，但仍抛出原 commit 异常。
- 重试策略
- run_in_transaction 在每次尝试内创建新的 TransactionContext，保证每次重试都是全新事务，不会污染状态。
- 使用 retriable_predicate 函数判断异常是否可重试（默认 is_postgres_retriable 识别 40001/40P01、类名或消息）。
- 指数回退+抖动避免惊群；deadline 可整体限制总时长。
- 可观测性
- on_event Hook 提供统一事件流（acquire.start/ok、begin.ok、commit/rollback、finish），易于接入Tracing/Metrics。
- 默认日志包含时序与异常上下文。
- 注意事项
- 关于驱动与数据库差异
- 代码中 SET TRANSACTION 与 SET LOCAL 语法适配 PostgreSQL。若使用 MySQL、MariaDB、SQLite 等，需要替换为等价指令或驱动API（例如 MySQL/MariaDB 可使用 SET SESSION 或在连接参数中指定超时）。
- autocommit 属性并非所有 DB-API 驱动都实现；实现中已做容错处理。
- 事务重试与幂等性
- 重试会导致数据库写操作多次尝试；请确保业务逻辑幂等（如使用唯一键、幂等键或去重表），避免重复插入。
- outbox 模式建议在同一事务内写入事件，异步发布并在消费者端保证至少一次投递与去重。
- 超时设置
- 语句级超时通过 SET LOCAL statement_timeout（毫秒）；这会在单条语句超时时由数据库中断该语句并抛出异常，应纳入可重试判定或异常处理。
- 连接池获取超时不等于数据库操作超时；两者分别配置。
- 资源清理
- 建议通过 session.cursor()/session.execute() 获取游标，便于上下文跟踪关闭。若直接使用 conn.cursor()，游标也会在 rollback/commit 后由服务器端回收，但客户端游标对象仍应关闭。
- 隔离级别选择
- 默认 REPEATABLE READ 在订单类场景通常合适；若存在高并发热点更新、写冲突频繁，可考虑 READ COMMITTED 并针对性加锁或采用乐观并发控制。若使用 SERIALIZABLE，请准备处理更高的序列化失败率并通过重试缓解。
- 异常不掩盖原则
- exit 在用户代码异常时始终返回 False，以便原始异常冒泡；回滚失败仅记录日志，绝不覆盖原异常。
- 生产建议
- 为 run_in_transaction 提供统一的 retry 策略与 on_event Hook，将日志、指标（尝试次数、失败原因、耗时分布）与追踪整合到平台（如 OpenTelemetry）。
- 对连接池大小、最大并发、语句超时进行容量规划与压测验证。
如需适配特定驱动（psycopg3、mysqlclient、aiomysql/asyncpg 等）或提供异步版本（async with），我可以在上述设计基础上给出对应的实现与最佳实践。
