Publish curated MobilityOps source snapshot

backend/app/core/db.py

@@ -0,0 +1,54 @@
+from collections.abc import Generator
+
+from sqlalchemy import create_engine, event, func, select
+from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
+
+from app.core.config import get_settings
+
+settings = get_settings()
+
+engine = create_engine(settings.database_url, pool_pre_ping=True, future=True)
+SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, future=True)
+
+# Every SQLAlchemy transaction participates in a shared database-wide barrier. Normal
+# reads/writes coexist; the short demo reset takes the exclusive form so it can never
+# interleave deletes/inserts with an API request, scanner, callback, or dispatcher cycle.
+DEMO_DATA_BARRIER_LOCK_ID = 5_344_725_149_212_793_901
+_EXCLUSIVE_RESET_INFO_KEY = "mobilityops_demo_reset_exclusive"
+
+
+@event.listens_for(Session, "after_begin")
+def _acquire_demo_data_barrier(session: Session, _transaction, connection) -> None:
+    if connection.dialect.name != "postgresql":
+        return
+    lock = (
+        func.pg_advisory_xact_lock(DEMO_DATA_BARRIER_LOCK_ID)
+        if session.info.get(_EXCLUSIVE_RESET_INFO_KEY)
+        else func.pg_advisory_xact_lock_shared(DEMO_DATA_BARRIER_LOCK_ID)
+    )
+    connection.execute(select(lock))
+
+
+def begin_exclusive_demo_reset(db: Session) -> None:
+    """Make the session's next transaction the exclusive side of the reset barrier."""
+    if db.in_transaction():
+        # Auth normally already read the user under a shared barrier. End that read-only
+        # transaction before requesting exclusive; in-place lock upgrades can deadlock.
+        db.rollback()
+    db.info[_EXCLUSIVE_RESET_INFO_KEY] = True
+
+
+def end_exclusive_demo_reset(db: Session) -> None:
+    db.info.pop(_EXCLUSIVE_RESET_INFO_KEY, None)
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+def get_db() -> Generator[Session, None, None]:
+    db = SessionLocal()
+    try:
+        yield db
+    finally:
+        db.close()