app/database.py

"""
database.py — Async PostgreSQL engine + session factory
"""

from __future__ import annotations

import logging
import uuid
from datetime import datetime
from typing import AsyncGenerator

from sqlalchemy import DateTime, String, func, text
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from app.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

# ── Engine ────────────────────────────────────────────────────────────────────
engine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    pool_pre_ping=True,
    pool_size=10,
    max_overflow=20,
    connect_args=settings.DB_CONNECT_ARGS,  # handles asyncpg SSL (Supabase/cloud)
)

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


# ── Base Model ────────────────────────────────────────────────────────────────
class Base(DeclarativeBase):
    """Common columns for every ORM model."""

    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
    )
    created_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
    )
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True),
        server_default=func.now(),
        onupdate=func.now(),
    )


# ── Dependency ────────────────────────────────────────────────────────────────
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI dependency – yields an async DB session per request."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


# ── Lifecycle helpers ─────────────────────────────────────────────────────────
async def create_tables() -> None:
    """Create all tables (used in startup when Alembic is not yet wired).

    copilot_chunks is created separately (see _create_copilot_table_if_possible)
    since it depends on the pgvector extension, which may not be installed/
    permitted on every deployment target — isolating it means a missing
    extension never blocks the rest of the schema from being created.
    """
    async with engine.begin() as conn:
        await conn.run_sync(
            Base.metadata.create_all,
            tables=[t for t in Base.metadata.sorted_tables if t.name != "copilot_chunks"],
        )
    await _create_copilot_table_if_possible()


async def _create_copilot_table_if_possible() -> None:
    try:
        from app.db.models.copilot import CopilotChunk

        async with engine.begin() as conn:
            await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
            await conn.run_sync(Base.metadata.create_all, tables=[CopilotChunk.__table__])
        logger.info("copilot_chunks table ready (pgvector enabled).")
    except Exception as exc:
        logger.warning(
            "pgvector unavailable (%s) — copilot chat/search will be disabled until it is.", exc
        )


async def drop_tables() -> None:
    """Drop all tables (testing only)."""
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.drop_all)