65 lines
1.7 KiB
Python
65 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import functools
|
|
import typing
|
|
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
)
|
|
from sqlalchemy.ext.asyncio import (
|
|
create_async_engine as sa_create_async_engine,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
|
|
from core.config import get_settings, is_development, is_test
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""Declarative base for all ORM models."""
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def get_engine() -> AsyncEngine:
|
|
"""Create and cache async engine from runtime settings."""
|
|
|
|
settings = get_settings()
|
|
return sa_create_async_engine(
|
|
settings.DATABASE_URL,
|
|
pool_size=settings.DB_POOL_SIZE,
|
|
max_overflow=settings.DB_MAX_OVERFLOW,
|
|
pool_timeout=settings.DB_POOL_TIMEOUT,
|
|
pool_recycle=settings.DB_POOL_RECYCLE,
|
|
)
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def get_session_factory() -> async_sessionmaker[AsyncSession]:
|
|
"""Create and cache async session factory."""
|
|
|
|
engine = get_engine()
|
|
return async_sessionmaker(engine, expire_on_commit=False)
|
|
|
|
|
|
AsyncSessionLocal = get_session_factory
|
|
|
|
|
|
async def get_db() -> typing.AsyncGenerator[AsyncSession, None]:
|
|
"""Dependency that yields an async database session."""
|
|
|
|
session_factory = get_session_factory()
|
|
async with session_factory() as session:
|
|
yield session
|
|
|
|
|
|
async def init_db_for_dev() -> None:
|
|
"""Initialize tables for development/test environment only."""
|
|
|
|
if not (is_development() or is_test()):
|
|
raise RuntimeError("init_db_for_dev can only be called in development or test")
|
|
|
|
engine = get_engine()
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(Base.metadata.create_all)
|