app/db/database.py

"""
Main application database (shared across all users).
  - users table
  - documents table (Qdrant vector registry)
  - conversations table (user chat sessions)
  - messages table (chat message history)
"""
import os
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy import String, Boolean, DateTime, Text, Integer, text
from datetime import datetime, timezone
from loguru import logger

from app.core.config import get_settings

settings = get_settings()

# Ensure data directory exists
os.makedirs("./data", exist_ok=True)


async def create_db_if_not_exists(database_url: str):
    # Only try to create database for postgresql URLs
    if not database_url.startswith("postgresql"):
        return
        
    from urllib.parse import urlparse
    parsed = urlparse(database_url)
    db_name = parsed.path.lstrip("/")
    
    # Connection URL to default 'postgres' database
    base_url = f"{parsed.scheme}://{parsed.netloc}/postgres"
    
    # Create a temporary engine connecting to the default 'postgres' DB outside of transaction
    temp_engine = create_async_engine(base_url, isolation_level="AUTOCOMMIT")
    try:
        async with temp_engine.connect() as conn:
            # Check if the database already exists
            result = await conn.execute(
                text("SELECT 1 FROM pg_database WHERE datname = :dbname"),
                {"dbname": db_name}
            )
            exists = result.scalar()
            if not exists:
                logger.info(f"Database '{db_name}' does not exist. Creating it...")
                await conn.execute(text(f"CREATE DATABASE {db_name}"))
                logger.info(f"Database '{db_name}' created successfully.")
    except Exception as e:
        logger.warning(f"Could not check/create database '{db_name}' automatically: {e}")
    finally:
        await temp_engine.dispose()


engine = create_async_engine(settings.database_url, echo=settings.debug)
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False)


class Base(DeclarativeBase):
    pass


class User(Base):
    __tablename__ = "users"

    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    email: Mapped[str] = mapped_column(String(255), unique=True, index=True, nullable=False)
    hashed_password: Mapped[str] = mapped_column(String(255), nullable=False)
    full_name: Mapped[str] = mapped_column(String(255), nullable=False)
    role: Mapped[str] = mapped_column(String(50), default="user")
    is_active: Mapped[bool] = mapped_column(Boolean, default=True)
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )


class DocumentRecord(Base):
    __tablename__ = "documents"

    id: Mapped[str] = mapped_column(String(36), primary_key=True)
    filename: Mapped[str] = mapped_column(String(512), nullable=False)
    title: Mapped[str] = mapped_column(String(512), nullable=False)
    source: Mapped[str | None] = mapped_column(String(512))
    chunk_count: Mapped[int] = mapped_column(Integer, default=0)
    uploaded_by: Mapped[str] = mapped_column(String(36), nullable=False)
    conversation_id: Mapped[str] = mapped_column(String(36), nullable=False, default="")
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
    )
    status: Mapped[str] = mapped_column(String(50), default="processing")
    error_message: Mapped[str | None] = mapped_column(Text)


async def init_db() -> None:
    # Auto-create the target PostgreSQL database if it does not exist
    await create_db_if_not_exists(settings.database_url)

    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        
        # Local import to prevent circular dependency
        from app.db.chat_db import ChatBase
        await conn.run_sync(ChatBase.metadata.create_all)
        
        # Use SQLAlchemy inspector to check columns database-agnostically
        def get_columns(connection):
            from sqlalchemy import inspect
            inspector = inspect(connection)
            return [col["name"] for col in inspector.get_columns("documents")]
            
        column_names = await conn.run_sync(get_columns)
        if "conversation_id" not in column_names:
            await conn.execute(text("ALTER TABLE documents ADD COLUMN conversation_id VARCHAR(36) NOT NULL DEFAULT ''"))


async def get_db():  # type: ignore[return]
    async with AsyncSessionLocal() as session:
        yield session