98 lines
2.9 KiB
Python
98 lines
2.9 KiB
Python
import logging
|
|
import os
|
|
from collections.abc import AsyncIterator
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from app.logging_utils import structured_log
|
|
from app.settings import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_engine: AsyncEngine | None = None
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
|
|
def _normalized_pool_mode(raw_mode: str) -> str:
|
|
mode = (raw_mode or "").strip().lower()
|
|
if mode == "auto":
|
|
if os.getenv("PYTEST_CURRENT_TEST"):
|
|
return "null"
|
|
app_env = (os.getenv("APP_ENV") or "").strip().lower()
|
|
if app_env in {"test", "development", "dev", "local"}:
|
|
return "null"
|
|
return "queue"
|
|
if mode in {"null", "queue"}:
|
|
return mode
|
|
structured_log(
|
|
logger,
|
|
"warning",
|
|
"db.invalid_pool_mode_fallback",
|
|
database_pool_mode=raw_mode,
|
|
fallback_mode="queue",
|
|
)
|
|
return "queue"
|
|
|
|
|
|
def get_engine() -> AsyncEngine:
|
|
global _engine
|
|
if _engine is None:
|
|
pool_mode = _normalized_pool_mode(settings.database_pool_mode)
|
|
engine_kwargs: dict[str, object] = {
|
|
"pool_pre_ping": True,
|
|
}
|
|
if pool_mode == "null":
|
|
engine_kwargs["poolclass"] = NullPool
|
|
else:
|
|
engine_kwargs["pool_size"] = max(1, int(settings.database_pool_size))
|
|
engine_kwargs["max_overflow"] = max(0, int(settings.database_pool_max_overflow))
|
|
engine_kwargs["pool_timeout"] = max(1, int(settings.database_pool_timeout_seconds))
|
|
|
|
_engine = create_async_engine(settings.database_url, **engine_kwargs)
|
|
structured_log(
|
|
logger,
|
|
"info",
|
|
"db.engine_initialized",
|
|
pool_mode=pool_mode,
|
|
)
|
|
return _engine
|
|
|
|
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession]:
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
_session_factory = async_sessionmaker(get_engine(), expire_on_commit=False)
|
|
return _session_factory
|
|
|
|
|
|
async def get_db_session() -> AsyncIterator[AsyncSession]:
|
|
session_factory = get_session_factory()
|
|
async with session_factory() as session:
|
|
yield session
|
|
|
|
|
|
async def check_database() -> bool:
|
|
engine = get_engine()
|
|
try:
|
|
async with engine.connect() as conn:
|
|
result = await conn.execute(text("SELECT 1"))
|
|
return result.scalar_one() == 1
|
|
except Exception:
|
|
structured_log(logger, "exception", "db.healthcheck_failed")
|
|
return False
|
|
|
|
|
|
async def close_engine() -> None:
|
|
global _engine, _session_factory
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
structured_log(logger, "info", "db.engine_disposed")
|
|
_engine = None
|
|
_session_factory = None
|