fix(db): enforce global nullpool session budget
This commit is contained in:
@@ -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:
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user