90 lines
2.7 KiB
Python
90 lines
2.7 KiB
Python
"""Database engine, session factory and schema bootstrap.
|
|
|
|
The engine is created lazily so importing the app (and loading the read-only
|
|
catalog) never requires a reachable database.
|
|
"""
|
|
import logging
|
|
import time
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.exc import IntegrityError, OperationalError, ProgrammingError
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.pool import StaticPool
|
|
|
|
from app import config
|
|
from app.models import Base
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
_engine = None
|
|
_session_factory = None
|
|
|
|
|
|
def get_engine():
|
|
"""Return the process-wide engine, creating it on first use."""
|
|
global _engine, _session_factory
|
|
if _engine is None:
|
|
url = config.DATABASE_URL
|
|
kwargs = {"pool_pre_ping": True, "future": True}
|
|
if url.startswith("sqlite"):
|
|
kwargs["connect_args"] = {"check_same_thread": False}
|
|
if ":memory:" in url or url in {"sqlite://", "sqlite+pysqlite://"}:
|
|
kwargs["poolclass"] = StaticPool
|
|
_engine = create_engine(url, **kwargs)
|
|
_session_factory = sessionmaker(
|
|
bind=_engine, autoflush=False, autocommit=False, expire_on_commit=False
|
|
)
|
|
return _engine
|
|
|
|
|
|
def get_session_factory():
|
|
get_engine()
|
|
return _session_factory
|
|
|
|
|
|
def get_db():
|
|
"""FastAPI dependency yielding a request-scoped SQLAlchemy session."""
|
|
db = get_session_factory()()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
|
|
def init_db(wait_seconds: int = 0, interval: float = 1.0) -> None:
|
|
"""Create missing tables, retrying while PostgreSQL finishes booting.
|
|
|
|
`create_all(checkfirst=True)` is idempotent and safe to run on every start.
|
|
It does not migrate changed columns; introduce Alembic when the schema
|
|
starts evolving in place.
|
|
"""
|
|
deadline = time.monotonic() + max(wait_seconds, 0)
|
|
while True:
|
|
try:
|
|
engine = get_engine()
|
|
with engine.connect() as conn:
|
|
conn.execute(text("SELECT 1"))
|
|
Base.metadata.create_all(bind=engine, checkfirst=True)
|
|
return
|
|
except (OperationalError, IntegrityError, ProgrammingError) as exc:
|
|
# A concurrent startup may have created the tables first.
|
|
if "already exists" in str(exc).lower():
|
|
return
|
|
if time.monotonic() >= deadline:
|
|
raise
|
|
log.warning(
|
|
"database not ready (%s); retrying in %.1fs",
|
|
exc.__class__.__name__,
|
|
interval,
|
|
)
|
|
time.sleep(interval)
|
|
|
|
|
|
def db_healthy() -> bool:
|
|
try:
|
|
with get_engine().connect() as conn:
|
|
conn.execute(text("SELECT 1"))
|
|
return True
|
|
except Exception:
|
|
return False
|