fix(db): enforce global nullpool session budget

This commit is contained in:
ogt
2026-07-15 15:49:18 +08:00
parent 19d7db93e5
commit f57845ba5e
12 changed files with 709 additions and 67 deletions

View File

@@ -83,14 +83,27 @@ class UnitOfWork:
"""進入事務"""
from src.core.context import get_current_project_id
self._session = self._session_factory()
if self._session is not None:
raise RuntimeError("UnitOfWork cannot be entered more than once")
session = self._session_factory()
self._session = session
effective_pid = (
self._project_id if self._project_id is not None else get_current_project_id()
)
await self._session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": effective_pid},
self._project_id
if self._project_id is not None
else get_current_project_id()
)
try:
await self._session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": effective_pid},
)
except BaseException:
# ``__aexit__`` is not invoked when ``__aenter__`` fails. Close
# here so a budgeted session cannot leak its global DB slot.
self._session = None
await session.close()
raise
self._committed = False
logger.debug("uow_started", project_id=effective_pid)
return self
@@ -109,22 +122,24 @@ class UnitOfWork:
- 無例外且未手動 commit: 自動 commit
- 已手動 commit: 不做任何事
"""
if exc_type is not None:
# 有例外rollback
await self.rollback()
logger.warning(
"uow_rollback_on_exception",
exc_type=exc_type.__name__ if exc_type else None,
exc_val=str(exc_val) if exc_val else None,
)
elif not self._committed:
# 無例外且未手動 commit自動 commit
await self.commit()
# 關閉 session
if self._session:
await self._session.close()
try:
if exc_type is not None:
# 有例外,rollback
await self.rollback()
logger.warning(
"uow_rollback_on_exception",
exc_type=exc_type.__name__ if exc_type else None,
exc_val=str(exc_val) if exc_val else None,
)
elif not self._committed:
# 無例外且未手動 commit自動 commit
await self.commit()
finally:
# Commit/rollback failures must also release a budgeted session.
session = self._session
self._session = None
if session:
await session.close()
async def commit(self) -> None:
"""