dependency.py

from typing import AsyncGenerator, Generator

from fastapi import Depends, FastAPI, Request, HTTPException, status
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.ext.asyncio import (
    create_async_engine,
    async_sessionmaker,
    AsyncSession,
    AsyncEngine,
)
from loguru import logger

from app.core import settings


def _async_engine_kwargs() -> dict:
    return {
        "echo": False,
        "pool_pre_ping": True,
        "pool_size": 5,
        "max_overflow": 10,
    }


def _sync_engine_kwargs() -> dict:
    return {
        "echo": False,
        "pool_pre_ping": True,
        "pool_size": 5,
        "max_overflow": 10,
    }


async def init_database_on_app(app: FastAPI) -> None:
    async_engine = create_async_engine(
        settings.ASYNC_DATABASE_URL, **_async_engine_kwargs()
    )
    sync_engine = create_engine(
        settings.SYNC_DATABASE_URL, **_sync_engine_kwargs()
    )

    app.state.async_engine = async_engine
    app.state.sync_engine = sync_engine
    app.state.async_sessionmaker = async_sessionmaker(
        bind=async_engine,
        expire_on_commit=False,
        autoflush=True,
    )
    app.state.sync_sessionmaker = sessionmaker(
        bind=sync_engine,
        expire_on_commit=False,
        autoflush=True,
    )
    logger.info("PostgreSQL engines & sessionmakers initialized.")


async def shutdown_database_on_app(app: FastAPI) -> None:
    if hasattr(app.state, "async_engine"):
        await app.state.async_engine.dispose()
    if hasattr(app.state, "sync_engine"):
        app.state.sync_engine.dispose()
    logger.info("PostgreSQL engines disposed.")


async def _ensure(
    app: FastAPI,
) -> tuple[async_sessionmaker[AsyncSession], sessionmaker]:
    if not hasattr(app.state, "async_sessionmaker"):
        logger.error("Async sessionmaker not initialized.")
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Service Unavailable",
        )
    if not hasattr(app.state, "sync_sessionmaker"):
        logger.error("Sync sessionmaker not initialized.")
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Service Unavailable",
        )
    return app.state.async_sessionmaker, app.state.sync_sessionmaker


async def _get_async_sessionmaker(
    request: Request,
) -> async_sessionmaker[AsyncSession]:
    async_maker, _ = await _ensure(request.app)
    return async_maker


async def _get_sync_sessionmaker(
    request: Request,
) -> sessionmaker:
    _, sync_maker = await _ensure(request.app)
    return sync_maker


async def get_async_session(
    session_maker: async_sessionmaker[AsyncSession] = Depends(
        _get_async_sessionmaker
    ),
) -> AsyncGenerator[AsyncSession, None]:
    async with session_maker() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


def get_sync_session(
    session_maker: sessionmaker = Depends(_get_sync_sessionmaker),
) -> Generator[Session, None, None]:
    session: Session = session_maker()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


_global_async_engine: AsyncEngine | None = None
_global_async_sessionmaker: async_sessionmaker[AsyncSession] | None = None


def _init_global_async_engine() -> AsyncEngine:
    global _global_async_engine
    if _global_async_engine is None:
        _global_async_engine = create_async_engine(
            settings.ASYNC_DATABASE_URL, **_async_engine_kwargs()
        )
        logger.info("Initialized global async engine.")
    return _global_async_engine


def _init_global_async_sessionmaker() -> async_sessionmaker[AsyncSession]:
    global _global_async_sessionmaker
    if _global_async_sessionmaker is None:
        engine = _init_global_async_engine()
        _global_async_sessionmaker = async_sessionmaker(
            bind=engine,
            expire_on_commit=False,
            autoflush=False,
            autocommit=False,
        )
        logger.info("Initialized global async sessionmaker.")
    return _global_async_sessionmaker


def get_global_async_sessionmaker() -> async_sessionmaker[AsyncSession]:
    return _init_global_async_sessionmaker()