python/config/database_config.py

"""
Database Configuration — PostgreSQL 15+ (asyncpg + SQLAlchemy async)

Handles:
- Async engine creation with connection pooling
- Session factory for dependency injection
- Database lifecycle (startup/shutdown)
- Health check utility
"""

import os
import logging
from sqlalchemy.ext.asyncio import (
    create_async_engine,
    AsyncSession,
    async_sessionmaker,
)
from sqlalchemy import text

# Base lives in models/ — database_config only handles engine & sessions
from models.base import Base

logger = logging.getLogger(__name__)


# ── Engine & Session Factory ──
DATABASE_URL = os.getenv(
    "DATABASE_URL",
    "postgresql+asyncpg://postgres:password@localhost:5432/apl_gs",  # pragma: allowlist secret
)

# ── Supabase Transaction Pooler Compatibility ──
# Supabase free tier only provides IPv4 via the Transaction Pooler (port 6543).
# PgBouncer in transaction mode does NOT support prepared statements.
# asyncpg uses prepared statements by default → must disable them.
#
# Detection: if the URL contains "pooler.supabase.com" or port 6543,
# we automatically add the fix.
_is_pooler = "pooler.supabase.com" in DATABASE_URL or ":6543/" in DATABASE_URL

_connect_args = {}
if _is_pooler:
    # Disable prepared statement cache for PgBouncer compatibility
    _connect_args["prepared_statement_cache_size"] = 0
    # Also disable statement_cache_size at the connection level
    _connect_args["statement_cache_size"] = 0
    logger.info("🔄 Supabase Transaction Pooler detected → disabled prepared statements")

engine = create_async_engine(
    DATABASE_URL,
    pool_size=int(os.getenv("DB_POOL_SIZE", "5" if _is_pooler else "10")),
    max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "10" if _is_pooler else "20")),
    pool_pre_ping=True,  # verify connections before use
    echo=os.getenv("APP_DEBUG", "False").lower() == "true",
    connect_args=_connect_args,
)

async_session_factory = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


# ── Dependency Injection ──
async def get_db() -> AsyncSession:
    """
    FastAPI dependency — yields an async database session.
    Automatically commits on success, rolls back on error, and closes.

    Usage:
        @router.get("/example")
        async def example(db: AsyncSession = Depends(get_db)):
            ...
    """
    async with async_session_factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


# ── Lifecycle ──
async def init_db() -> None:
    """
    Create all tables defined by ORM models.
    Called once at application startup.

    All ORM models must be imported before create_all so SQLAlchemy
    knows about every table. They are imported here explicitly.
    """
    # Ensure all model classes are registered with Base.metadata
    import models.outlet_model  # noqa: F401
    import models.visit_model  # noqa: F401
    import models.checklist_model  # noqa: F401
    import models.photo_model  # noqa: F401
    import models.assessment_model  # noqa: F401

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    logger.info("✅ Database tables created / verified")


async def close_db() -> None:
    """Dispose the engine connection pool on shutdown."""
    await engine.dispose()
    logger.info("🔌 Database connection pool closed")


# ── Health Check ──
async def check_db_health() -> dict:
    """
    Returns database connectivity status.
    Used by the /api/v1/health endpoint.
    """
    try:
        async with async_session_factory() as session:
            result = await session.execute(text("SELECT 1"))
            result.scalar()
        return {"database": "connected", "status": "healthy"}
    except Exception as e:
        logger.error(f"Database health check failed: {e}")
        return {"database": "disconnected", "status": "unhealthy", "error": str(e)}