File size: 3,326 Bytes
10795db 34cb656 10795db 34cb656 10795db 34cb656 10795db 34cb656 10795db | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from collections.abc import Generator
from sqlalchemy import create_engine, event
from sqlalchemy.engine import Engine
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
from sqlalchemy.pool import StaticPool
from core.config import settings
_is_sqlite = settings.DATABASE_URL.startswith("sqlite")
_is_memory = _is_sqlite and (":memory:" in settings.DATABASE_URL or "mode=memory" in settings.DATABASE_URL)
_connect_args = {"check_same_thread": False} if _is_sqlite else {}
def configure_sqlite(bound_engine: Engine, *, in_memory: bool | None = None) -> None:
"""Apply the SQLite pragmas the concurrent ingest path relies on.
Alert-time Jarvis enrichment now commits from a worker thread while the
ingest coroutine (and other live requests) hold their own connections, so
two writers can race for the database lock. Without a busy timeout SQLite
fails that race immediately with "database is locked"; with one, the loser
waits its (brief) turn instead. WAL additionally keeps a write from
blocking dashboard reads — but it is only valid on an on-disk database, so
we skip it for in-memory test databases.
"""
if in_memory is None:
url = str(bound_engine.url)
in_memory = ":memory:" in url or "mode=memory" in url
@event.listens_for(bound_engine, "connect")
def _set_sqlite_pragmas(dbapi_conn, _record): # pragma: no cover - trivial
cursor = dbapi_conn.cursor()
cursor.execute("PRAGMA busy_timeout=5000")
if not in_memory:
cursor.execute("PRAGMA journal_mode=WAL")
cursor.execute("PRAGMA synchronous=NORMAL")
cursor.close()
_engine_kwargs: dict = {
"connect_args": _connect_args,
"echo": settings.DEBUG,
"future": True,
}
if _is_memory:
# A shared in-memory SQLite database exists only while a connection to it
# is open. Now that alert-time enrichment opens its own session from a
# worker thread, a QueuePool could momentarily hold zero connections and
# destroy the database mid-run. StaticPool keeps a single long-lived
# connection so the schema survives across threads (SQLite serializes
# access internally). On-disk deployments keep the default pool.
_engine_kwargs["poolclass"] = StaticPool
engine = create_engine(settings.DATABASE_URL, **_engine_kwargs)
if _is_sqlite:
configure_sqlite(engine, in_memory=_is_memory)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, class_=Session)
class Base(DeclarativeBase):
pass
def get_db() -> Generator[Session, None, None]:
db = SessionLocal()
try:
yield db
finally:
db.close()
def init_db() -> None:
# Importing here ensures every model is registered against Base.metadata
# before create_all walks the registry.
from models import Base as ModelsBase # noqa: F401
Base.metadata.create_all(bind=engine)
# Auto-add columns introduced after the demo DB was first created.
# Safe + idempotent — no-op when the schema is already up to date.
try:
from core.migrations import upgrade_schema
upgrade_schema(engine)
except Exception: # pragma: no cover - never block startup
import logging
logging.getLogger("biopatch").exception("Auto-migration failed")
|