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