backend/app/core/database.py

"""
Async Database Setup
Supports PostgreSQL (local dev) and SQLite (HF Spaces) via DATABASE_URL.
"""
from __future__ import annotations

import os
import structlog
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase

from app.core.config import settings

logger = structlog.get_logger(__name__)

# ── Engine ───────────────────────────────────────────────────
connect_args = {}
if settings.use_sqlite:
    connect_args["check_same_thread"] = False

engine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    connect_args=connect_args,
)

AsyncSessionLocal = async_sessionmaker(
    bind=engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)


class Base(DeclarativeBase):
    """Base class for all SQLAlchemy models."""
    pass


async def init_db() -> None:
    """Create all tables on startup (ensures all tables exist in PG & SQLite)."""
    # Import all models here so Base knows about all tables
    from app.models import tenant, user, audit, document, chat  # noqa: F401

    if settings.use_sqlite:
        os.makedirs("./data", exist_ok=True)

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
    logger.info("Database tables initialized successfully", db=settings.DATABASE_URL)


async def get_db() -> AsyncSession:
    """FastAPI dependency: provides async DB session."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()