81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
from contextlib import contextmanager
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, DeclarativeBase
|
|
|
|
from config import settings
|
|
|
|
|
|
def _pool_kwargs(url: str) -> dict:
|
|
"""QueuePool settings for non-SQLite engines.
|
|
|
|
SQLite keeps its own default pool (SingletonThreadPool for file DBs) which
|
|
does not accept QueuePool kwargs — and has no pool exhaustion problem since
|
|
it is single-connection-per-thread. Values are clamped so a bad env value
|
|
can never shrink the pool to zero or make the timeout negative.
|
|
"""
|
|
if url.startswith("sqlite"):
|
|
return {}
|
|
return {
|
|
"pool_size": max(1, int(settings.db_pool_size)),
|
|
"max_overflow": max(0, int(settings.db_max_overflow)),
|
|
"pool_timeout": max(0.0, float(settings.db_pool_timeout)),
|
|
"pool_recycle": max(0, int(settings.db_pool_recycle)),
|
|
"pool_use_lifo": bool(settings.db_pool_use_lifo),
|
|
}
|
|
|
|
|
|
def build_engine(url: str | None = None, **kwargs):
|
|
"""Create the app engine with the configured pool settings.
|
|
|
|
``url`` defaults to settings.database_url. Extra kwargs (e.g. pool_size for
|
|
tests) override the configured values. `check_same_thread=False` is applied
|
|
only for SQLite (FastAPI serves requests from a thread pool); passing it to
|
|
the PostgreSQL driver (psycopg2) raises TypeError.
|
|
"""
|
|
url = url or settings.database_url
|
|
connect_args = {"check_same_thread": False} if url.startswith("sqlite") else {}
|
|
return create_engine(
|
|
url,
|
|
connect_args=connect_args,
|
|
# Ping the connection before checkout so idle PostgreSQL connections that
|
|
# were reaped by the server get transparently replaced instead of raising
|
|
# "server closed the connection unexpectedly". No-op for SQLite.
|
|
pool_pre_ping=True,
|
|
**_pool_kwargs(url),
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
engine = build_engine()
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
@contextmanager
|
|
def db_session():
|
|
"""Short-lived session for background jobs and scripts.
|
|
|
|
Yields a Session that is ALWAYS closed (rolling back any uncommitted
|
|
transaction) when the block exits — the job must not hold a Session across
|
|
long external awaits (Gemini/embedding/HTTP calls) because the session's
|
|
open transaction pins a pooled connection for the whole call. Prefer this
|
|
over ``db = SessionLocal()`` in long-running code.
|
|
"""
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|