api/database.py

"""
Database connection and session management
"""
import os
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import StaticPool
from typing import Generator

from api.models import Base

# Database URL priority:
# 1. NEON_DATABASE_URL_DEV (local development with PostgreSQL)
# 2. NEON_DATABASE_URL (production Neon PostgreSQL)
# 3. DATABASE_URL (backwards compatibility)
# 4. SQLite fallback (no setup required)
DATABASE_URL = os.getenv(
    "NEON_DATABASE_URL_DEV",
    os.getenv(
        "NEON_DATABASE_URL",
        os.getenv(
            "DATABASE_URL",
            "sqlite:///./data/users.db"
        )
    )
)

# Handle PostgreSQL URL format for SQLAlchemy 2.0+
if DATABASE_URL.startswith("postgres://"):
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)

# Create engine
if "sqlite" in DATABASE_URL:
    # SQLite needs special handling for concurrent access
    engine = create_engine(
        DATABASE_URL,
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
else:
    # PostgreSQL configuration
    engine = create_engine(
        DATABASE_URL,
        pool_pre_ping=True,
        pool_size=10,
        max_overflow=20,
    )

# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


def init_db():
    """Create all tables"""
    Base.metadata.create_all(bind=engine)
    print(f"✅ Database initialized at: {DATABASE_URL}")


def get_db() -> Generator[Session, None, None]:
    """
    Database session dependency for FastAPI
    
    Usage:
        @app.get("/users")
        def get_users(db: Session = Depends(get_db)):
            return db.query(User).all()
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()