punimtag/backend/db/session.py

from __future__ import annotations

from pathlib import Path
from typing import Generator

from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

# Load environment variables from .env file if it exists
# Path: backend/db/session.py -> backend/db -> backend -> punimtag/ -> .env
env_path = Path(__file__).parent.parent.parent / ".env"
load_dotenv(dotenv_path=env_path)


def get_database_url() -> str:
    """Fetch database URL from environment or defaults."""
    import os
    # Check for environment variable first
    db_url = os.getenv("DATABASE_URL")
    if db_url:
        return db_url
    # Default to PostgreSQL for development (without password - must be set via env var)
    # This ensures no hardcoded passwords in the codebase
    raise ValueError(
        "DATABASE_URL environment variable not set. "
        "Please set DATABASE_URL in your .env file or environment."
    )


def get_auth_database_url() -> str:
    """Fetch auth database URL from environment."""
    import os
    db_url = os.getenv("DATABASE_URL_AUTH")
    if not db_url:
        raise ValueError("DATABASE_URL_AUTH environment variable not set")
    return db_url


database_url = get_database_url()
# PostgreSQL connection pool settings
pool_kwargs = {
    "pool_pre_ping": True,
    "pool_size": 10,
    "max_overflow": 20,
    "pool_recycle": 3600,
}

engine = create_engine(
    database_url,
    future=True,
    **pool_kwargs
)
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)


def get_db() -> Generator:
    """Yield a DB session for request lifecycle."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# Auth database setup
try:
    auth_database_url = get_auth_database_url()
    auth_pool_kwargs = {
        "pool_pre_ping": True,
        "pool_size": 10,
        "max_overflow": 20,
        "pool_recycle": 3600,
    }
    
    auth_engine = create_engine(
        auth_database_url,
        future=True,
        **auth_pool_kwargs
    )
    AuthSessionLocal = sessionmaker(bind=auth_engine, autoflush=False, autocommit=False, future=True)
except ValueError as e:
    # DATABASE_URL_AUTH not set - auth database not available
    print(f"[DB Session] ⚠️  Auth database not configured: {e}")
    auth_engine = None
    AuthSessionLocal = None
except Exception as e:
    # Other errors (connection failures, etc.) - log but don't crash
    import os
    print(f"[DB Session] ⚠️  Failed to initialize auth database: {e}")
    print(f"[DB Session]    URL was: {os.getenv('DATABASE_URL_AUTH', 'not set')}")
    auth_engine = None
    AuthSessionLocal = None


def get_auth_db() -> Generator:
    """Yield a DB session for auth database request lifecycle."""
    if AuthSessionLocal is None:
        from fastapi import HTTPException, status
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Auth database not configured. Please set DATABASE_URL_AUTH environment variable in the backend configuration."
        )
    db = AuthSessionLocal()
    try:
        yield db
    finally:
        db.close()