55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
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()
|