Spaces:
Sleeping
Sleeping
File size: 4,064 Bytes
2dca8fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """
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)
|