Spaces:
Runtime error
Runtime error
File size: 1,795 Bytes
31e8a9f | 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 | """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)
@contextmanager
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()
|