app/core/database.py

"""
Database connection module with async support.
"""
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker, declarative_base

from app.core.config import settings

# Get the database URL from settings
db_url = str(settings.DATABASE_URL)

# For SQLite, don't add asyncpg prefix
if db_url.startswith("sqlite"):
    engine = create_async_engine(
        db_url,
        echo=settings.DEBUG,
    )
else:
    # For PostgreSQL, add asyncpg prefix
    engine = create_async_engine(
        db_url.replace("postgresql://", "postgresql+asyncpg://"),
        pool_pre_ping=True,
        echo=settings.DEBUG,
    )

# Create async SessionLocal class
AsyncSessionLocal = sessionmaker(
    engine, class_=AsyncSession, expire_on_commit=False, autocommit=False, autoflush=False
)

# Backwards compatibility alias
SessionLocal = AsyncSessionLocal

# Create Base class for models
Base = declarative_base()


async def get_db() -> AsyncSession:
    """
    Dependency to get async DB session.
    """
    async with AsyncSessionLocal() as session:
        yield session