core-api/src/db/database.py

"""
Database Connection Module

Provides async PostgreSQL connectivity using SQLAlchemy 2.0 with asyncpg driver.
Follows the existing singleton pattern used throughout core-api.
"""
from typing import AsyncGenerator, Optional
from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncSession,
    AsyncEngine,
    create_async_engine,
    async_sessionmaker,
)
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.pool import NullPool

from src.config import get_settings
from src.logging_config import get_logger

logger = get_logger(__name__)
settings = get_settings()


class Base(DeclarativeBase):
    """
    SQLAlchemy declarative base for all models

    All database models should inherit from this class.
    """
    pass


class Database:
    """
    Async database connection manager

    Provides async engine and session factory for PostgreSQL connections.
    Uses asyncpg driver for optimal async performance.
    """

    def __init__(self, database_url: Optional[str] = None):
        """
        Initialize database connection manager

        Args:
            database_url: PostgreSQL connection URL (default from settings)
        """
        # Convert postgresql:// to postgresql+asyncpg:// for async driver
        url = database_url or settings.database_url
        if url.startswith("postgresql://"):
            url = url.replace("postgresql://", "postgresql+asyncpg://", 1)

        self._url = url
        self._engine: Optional[AsyncEngine] = None
        self._session_factory: Optional[async_sessionmaker[AsyncSession]] = None

    @property
    def engine(self) -> AsyncEngine:
        """
        Get or create the async database engine

        Returns:
            AsyncEngine instance
        """
        if self._engine is None:
            self._engine = create_async_engine(
                self._url,
                echo=settings.debug,  # Log SQL in debug mode
                poolclass=NullPool,  # Disable connection pooling for serverless compatibility
            )
            logger.info(f"Database engine created for {self._url.split('@')[-1]}")
        return self._engine

    @property
    def session_factory(self) -> async_sessionmaker[AsyncSession]:
        """
        Get or create the async session factory

        Returns:
            Session factory for creating database sessions
        """
        if self._session_factory is None:
            self._session_factory = async_sessionmaker(
                bind=self.engine,
                class_=AsyncSession,
                expire_on_commit=False,
                autocommit=False,
                autoflush=False,
            )
        return self._session_factory

    async def create_tables(self) -> None:
        """
        Create all database tables

        Should only be used for development/testing.
        Use Alembic migrations for production.
        """
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
        logger.info("Database tables created")

    async def drop_tables(self) -> None:
        """
        Drop all database tables

        WARNING: Destroys all data. Use with caution.
        """
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)
        logger.warning("Database tables dropped")

    async def health_check(self) -> bool:
        """
        Check if database connection is healthy

        Returns:
            True if connection successful, False otherwise
        """
        try:
            async with self.session_factory() as session:
                await session.execute(text("SELECT 1"))
            return True
        except Exception as e:
            logger.error(f"Database health check failed: {e}")
            return False

    async def close(self) -> None:
        """
        Close database connections and dispose of engine
        """
        if self._engine is not None:
            await self._engine.dispose()
            self._engine = None
            self._session_factory = None
            logger.info("Database connections closed")


# Singleton instance
_database: Optional[Database] = None


def get_database() -> Database:
    """
    Get singleton database instance

    Returns:
        Database instance
    """
    global _database
    if _database is None:
        _database = Database()
    return _database


async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
    """
    FastAPI dependency for database sessions

    Yields an async session that is automatically closed after the request.

    Usage:
        @router.get("/items")
        async def get_items(session: AsyncSession = Depends(get_async_session)):
            result = await session.execute(select(Item))
            return result.scalars().all()

    Yields:
        AsyncSession instance
    """
    database = get_database()
    async with database.session_factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise