Spaces:
Runtime error
Runtime error
File size: 2,165 Bytes
00a4363 34b0226 00a4363 34b0226 66de9c8 4a3d0a5 66de9c8 34b0226 00a4363 | 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 | from collections.abc import Iterator
from sqlalchemy import create_engine, event
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
from app.config import settings
class Base(DeclarativeBase):
pass
_sqlite = settings.database_url.startswith("sqlite")
# A hosted Postgres (Neon) closes idle connections and drops the odd one on a network blip, and a connection that dies silently
# (laptop sleep, network change) otherwise leaves a query waiting on the socket forever, and the worker loop with it. TCP
# keepalives notice a dead peer within about a minute, so the query fails instead and a dead socket is noticed while it sits in
# the pool; connect_timeout does the same for a connect that never completes. A pooled connection is also replaced before it
# is handed out (pool_recycle, below Neon's idle timeout), since pool_pre_ping alone can fail on a dead SSL socket with a
# ProgrammingError ("can't change 'autocommit' now") that SQLAlchemy does not treat as a lost connection, and the request
# ends in a 500.
_pg_connect_args = {"connect_timeout": 10, "keepalives": 1, "keepalives_idle": 30, "keepalives_interval": 10,
"keepalives_count": 3}
_connect_args = {"check_same_thread": False, "timeout": 30} if _sqlite else _pg_connect_args
engine = create_engine(settings.database_url, pool_pre_ping=True, connect_args=_connect_args,
**({} if _sqlite else {"pool_recycle": 240, "pool_use_lifo": True}))
if _sqlite:
# Local dev: the API and the worker share one file. WAL lets reads run during a write, the busy
# timeout makes a second writer wait instead of failing, and SQLite needs FKs switched on.
@event.listens_for(engine, "connect")
def _sqlite_pragmas(dbapi_conn, _record):
cur = dbapi_conn.cursor()
cur.execute("PRAGMA journal_mode=WAL")
cur.execute("PRAGMA busy_timeout=30000")
cur.execute("PRAGMA foreign_keys=ON")
cur.close()
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
def get_db() -> Iterator[Session]:
db = SessionLocal()
try:
yield db
finally:
db.close()
|