Spaces:
Runtime error
Runtime error
| """SQLAlchemy engine, session factory, and declarative base.""" | |
| from __future__ import annotations | |
| from collections.abc import Iterator | |
| from contextlib import contextmanager | |
| from sqlalchemy import JSON, create_engine | |
| from sqlalchemy.dialects.postgresql import JSONB | |
| from sqlalchemy.engine import Engine | |
| from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker | |
| from sqlalchemy.pool import StaticPool | |
| from app.config import settings | |
| #: JSON column type: native JSONB on Postgres, plain JSON elsewhere (SQLite tests/CI). | |
| JSONType = JSON().with_variant(JSONB(), "postgresql") | |
| class Base(DeclarativeBase): | |
| """Declarative base for all ORM models.""" | |
| def build_engine(url: str) -> Engine: | |
| """Create an engine with settings appropriate for the URL's dialect.""" | |
| if url.startswith("sqlite"): | |
| # In-memory SQLite needs a single shared connection across threads. | |
| in_memory = ":memory:" in url or url.rstrip("/") == "sqlite:" | |
| return create_engine( | |
| url, | |
| connect_args={"check_same_thread": False}, | |
| poolclass=StaticPool if in_memory else None, | |
| ) | |
| return create_engine( | |
| url, | |
| pool_pre_ping=True, # survive Supabase/pooler idle disconnects | |
| pool_size=5, | |
| max_overflow=5, | |
| pool_recycle=1800, | |
| ) | |
| engine = build_engine(settings.DATABASE_URL) | |
| SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False) | |
| def session_scope() -> Iterator[Session]: | |
| """Context-managed session for non-request code (Celery tasks, scripts). | |
| Rolls back on any exception and always closes the session. | |
| """ | |
| db = SessionLocal() | |
| try: | |
| yield db | |
| except Exception: | |
| db.rollback() | |
| raise | |
| finally: | |
| db.close() | |