""" 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)