wooo/awoooi
1
1
Fork 0
Code Issues 2 Pull Requests 1 Actions 2 Packages Projects Releases Wiki Activity

fix(db): enforce global nullpool session budget

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

57
apps/api/src/core/unit_of_work.py
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:
"""

222
apps/api/src/db/base.py
View File

@@ -16,6 +16,7 @@ Features:
import asyncio
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from typing import Any
from fastapi import HTTPException
from sqlalchemy import text
@@ -25,6 +26,7 @@ from sqlalchemy.ext.asyncio import (
AsyncConnection,
AsyncEngine,
AsyncSession,
AsyncSessionTransaction,
async_sessionmaker,
create_async_engine,
)
@@ -98,6 +100,160 @@ async def _database_session_slot() -> AsyncGenerator[None, None]:
limiter.release()
class _BudgetedAsyncSession(AsyncSession):
"""Apply one NullPool concurrency lease to every session access path.
``get_db`` and ``get_db_context`` used to own the semaphore lease outside
the SQLAlchemy session. Callers that consumed ``get_session_factory()``
directly (notably ``UnitOfWork``) bypassed that boundary entirely. Keep
the lease on the session instead so dependency, context-manager, raw
factory, and transaction APIs all share the same per-process budget.
The lease is lazy for sessions constructed but never used, and is held
until close/reset/invalidate so a single session cannot churn through the
semaphore while retaining transactional state.
"""
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._database_slot_context: Any | None = None
self._database_slot_lock = asyncio.Lock()
async def _acquire_database_slot(self) -> None:
if self._database_slot_context is not None:
return
async with self._database_slot_lock:
if self._database_slot_context is not None:
return
slot_context = _database_session_slot()
await slot_context.__aenter__()
self._database_slot_context = slot_context
async def _release_database_slot(self) -> None:
async with self._database_slot_lock:
slot_context = self._database_slot_context
self._database_slot_context = None
if slot_context is not None:
await slot_context.__aexit__(None, None, None)
async def __aenter__(self) -> "_BudgetedAsyncSession":
slot_was_held = self._database_slot_context is not None
await self._acquire_database_slot()
try:
return await super().__aenter__()
except BaseException:
# A failed context entry has no matching ``__aexit__``. Close a
# lease acquired by this entry so cancellation/future SQLAlchemy
# entry failures cannot strand one of the process-wide slots.
if not slot_was_held:
await self.close()
raise
def begin(self) -> AsyncSessionTransaction:
return _BudgetedAsyncSessionTransaction(self)
def begin_nested(self) -> AsyncSessionTransaction:
return _BudgetedAsyncSessionTransaction(self, nested=True)
async def connection(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().connection(*args, **kwargs)
async def delete(self, *args: Any, **kwargs: Any) -> None:
await self._acquire_database_slot()
await super().delete(*args, **kwargs)
async def execute(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().execute(*args, **kwargs)
async def flush(self, *args: Any, **kwargs: Any) -> None:
await self._acquire_database_slot()
await super().flush(*args, **kwargs)
async def get(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().get(*args, **kwargs)
async def get_one(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().get_one(*args, **kwargs)
async def merge(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().merge(*args, **kwargs)
async def refresh(self, *args: Any, **kwargs: Any) -> None:
await self._acquire_database_slot()
await super().refresh(*args, **kwargs)
async def run_sync(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().run_sync(*args, **kwargs)
async def scalar(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().scalar(*args, **kwargs)
async def scalars(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().scalars(*args, **kwargs)
async def stream(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().stream(*args, **kwargs)
async def stream_scalars(self, *args: Any, **kwargs: Any) -> Any:
await self._acquire_database_slot()
return await super().stream_scalars(*args, **kwargs)
async def commit(self) -> None:
await self._acquire_database_slot()
await super().commit()
async def close(self) -> None:
try:
await super().close()
finally:
await self._release_database_slot()
async def reset(self) -> None:
try:
await super().reset()
finally:
await self._release_database_slot()
async def invalidate(self) -> None:
try:
await super().invalidate()
finally:
await self._release_database_slot()
class _BudgetedAsyncSessionTransaction(AsyncSessionTransaction):
"""Acquire the session lease for explicit ``begin`` entry paths."""
async def start(
self,
is_ctxmanager: bool = False,
) -> AsyncSessionTransaction:
session = self.session
if isinstance(session, _BudgetedAsyncSession):
slot_was_held = session._database_slot_context is not None
await session._acquire_database_slot()
try:
return await super().start(is_ctxmanager=is_ctxmanager)
except BaseException:
# ``async_sessionmaker.begin()`` cannot call its exit path
# when transaction entry fails. Close only when this start
# acquired the lease; an already-entered session still owns
# its pre-existing lease and lifecycle.
if not slot_was_held:
await session.close()
raise
return await super().start(is_ctxmanager=is_ctxmanager)
def _raise_unauthorized_db_context(msg: str) -> None:
context = get_current_project_context()
logger.error(
@@ -152,12 +308,12 @@ def get_engine() -> AsyncEngine:
def get_session_factory() -> async_sessionmaker[AsyncSession]:
"""Get or create session factory"""
"""Get the globally budgeted async session factory."""
global _session_factory
if _session_factory is None:
_session_factory = async_sessionmaker(
bind=get_engine(),
class_=AsyncSession,
class_=_BudgetedAsyncSession,
expire_on_commit=False,
autoflush=False,
)
@@ -178,28 +334,27 @@ async def get_db() -> AsyncGenerator[AsyncSession, None]:
...
"""
factory = get_session_factory()
async with _database_session_slot():
async with factory() as session:
try:
from src.core.context import get_current_project_id
async with factory() as session:
try:
from src.core.context import get_current_project_id
# AwoooP Phase 2.3 (2026-05-04 ogt): SET LOCAL app.project_id 讓 RLS Policy 生效
# Fail-Closed RLS: 遇到未授權情境拋出錯誤而非回退到 "awoooi"
pid = get_current_project_id()
if not pid:
_raise_unauthorized_db_context(
"Unauthorized: project_id is missing in context (Fail-Closed RLS)"
)
await session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": pid},
# AwoooP Phase 2.3 (2026-05-04 ogt): SET LOCAL app.project_id 讓 RLS Policy 生效
# Fail-Closed RLS: 遇到未授權情境拋出錯誤而非回退到 "awoooi"
pid = get_current_project_id()
if not pid:
_raise_unauthorized_db_context(
"Unauthorized: project_id is missing in context (Fail-Closed RLS)"
)
yield session
await session.commit()
except Exception:
await session.rollback()
raise
await session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": pid},
)
yield session
await session.commit()
except Exception:
await session.rollback()
raise
@asynccontextmanager
@@ -224,18 +379,17 @@ async def get_db_context(project_id: str | None = None) -> AsyncGenerator[AsyncS
_raise_unauthorized_db_context("Unauthorized: project_id is missing in context (Fail-Closed RLS)")
factory = get_session_factory()
async with _database_session_slot():
async with factory() as session:
try:
await session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": effective_pid},
)
yield session
await session.commit()
except Exception:
await session.rollback()
raise
async with factory() as session:
try:
await session.execute(
text("SELECT set_config('app.project_id', :pid, TRUE)"),
{"pid": effective_pid},
)
yield session
await session.commit()
except Exception:
await session.rollback()
raise
# =============================================================================

14
apps/api/src/services/executor_trust_boundary_readback.py
View File

@@ -154,6 +154,11 @@ def _build_database_identity_boundary(
"database_null_pool": _truthy(
receipt.get(f"{workload_id}_database_null_pool")
),
"null_pool_concurrency_limit": _int_or_none(
receipt.get(
f"{workload_id}_database_null_pool_concurrency_limit"
)
),
"monolithic_secret_env_from_absent": _truthy(
receipt.get(
f"{workload_id}_monolithic_secret_env_from_absent"
@@ -198,6 +203,9 @@ def _build_database_identity_boundary(
workloads[workload_id]["database_null_pool"] is True
for workload_id in _DB_WORKLOAD_IDS
)
api_null_pool_concurrency_limit_verified = (
workloads["api"]["null_pool_concurrency_limit"] == 2
)
workload_budget_verified = all(
workloads[workload_id]["connection_limit"]
== _DB_WORKLOAD_BUDGETS[workload_id]["connection_limit"]
@@ -256,6 +264,7 @@ def _build_database_identity_boundary(
connection_budget_verified = bool(
db_identity_isolated
and null_pool_verified
and api_null_pool_concurrency_limit_verified
and workload_budget_verified
and representative_preflights_verified
and verifier_verified
@@ -276,6 +285,8 @@ def _build_database_identity_boundary(
blockers.append("database_role_fingerprints_not_distinct")
if not api_http_concurrency_verified:
blockers.append("api_http_concurrency_probe_not_verified")
if not api_null_pool_concurrency_limit_verified:
blockers.append("api_null_pool_concurrency_limit_not_verified")
if not verifier_verified:
blockers.append("database_identity_verifier_not_verified")
if verified_at_utc is None:
@@ -360,6 +371,9 @@ def _build_database_identity_boundary(
"distinct_database_secret_refs": distinct_secret_refs,
"distinct_db_role_fingerprints": distinct_role_fingerprints,
"database_null_pool_verified": null_pool_verified,
"api_null_pool_concurrency_limit_verified": (
api_null_pool_concurrency_limit_verified
),
"workload_connection_budgets_verified": workload_budget_verified,
"database_identity_verifier_verified": verifier_verified,
"database_capacity_receipt_fresh": capacity_receipt_fresh,