HumeAI/expressive-tts-arena

1# Standard Library Imports2from typing import Callable, Optional, TypeAlias, Union3 4# Third-Party Library Imports5from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine6from sqlalchemy.orm import DeclarativeBase7 8# Local Application Imports9from src.common import Config, logger10 11 12# Define the SQLAlchemy Base13class Base(DeclarativeBase):14    pass15 16class DummyAsyncSession:17    is_dummy = True  # Flag to indicate this is a dummy session.18 19    async def __enter__(self):20        return self21 22    async def __exit__(self, exc_type, exc_value, traceback):23        pass24 25    async def add(self, _instance, _warn=True):26        # No-op: simply ignore adding the instance.27        pass28 29    async def commit(self):30        # Raise an exception to simulate failure when attempting a write.31        raise RuntimeError("DummyAsyncSession does not support commit operations.")32 33    async def refresh(self, _instance):34        # Raise an exception to simulate failure when attempting to refresh.35        raise RuntimeError("DummyAsyncSession does not support refresh operations.")36 37    async def rollback(self):38        # No-op: there's nothing to roll back.39        pass40 41    async def close(self):42        # No-op: nothing to close.43        pass44 45AsyncDBSessionMaker: TypeAlias = Union[async_sessionmaker[AsyncSession], Callable[[], DummyAsyncSession]]46engine: Optional[AsyncEngine] = None47 48def init_db(config: Config) -> AsyncDBSessionMaker:49    """50    Initialize the database engine and return a session factory based on the provided configuration.51 52    In production, a valid DATABASE_URL is required. In development, if a DATABASE_URL is provided,53    it is used; otherwise, a dummy session factory is returned to allow graceful failure.54 55    Args:56        config (Config): The application configuration.57 58    Returns:59        AsyncDBSessionMaker: A sessionmaker bound to the engine, or a dummy session factory.60    """61    # ruff doesn't like setting global variables, but this is practical here62    global engine  # noqa63 64    if config.app_env == "prod":65        # In production, a valid DATABASE_URL is required.66        if not config.database_url:67            raise ValueError("DATABASE_URL must be set in production!")68 69        async_db_url = config.database_url70        engine = create_async_engine(async_db_url)71 72        return async_sessionmaker(bind=engine, expire_on_commit=False, class_=AsyncSession)73 74    # In development, if a DATABASE_URL is provided, use it.75    if config.database_url:76        async_db_url = config.database_url77        engine = create_async_engine(async_db_url)78 79        return async_sessionmaker(bind=engine, expire_on_commit=False, class_=AsyncSession)80 81    # No DATABASE_URL is provided; use a DummyAsyncSession that does nothing.82    engine = None83    logger.warning("No DATABASE_URL provided - database operations will use DummyAsyncSession")84 85    def async_dummy_session_factory() -> DummyAsyncSession:86        return DummyAsyncSession()87 88    return async_dummy_session_factory89