tradedangerous 12.7.6__py3-none-any.whl

tradedangerous/db/locks.py

@@ -0,0 +1,208 @@
+# tradedangerous/db/locks.py
+# -----------------------------------------------------------------------------
+# Advisory lock helpers (MariaDB/MySQL) — per-station serialization
+#
+# SQLite compatibility:
+#   - On SQLite (or any unsupported dialect), all helpers become NO-OPs and
+#     behave as if the lock was immediately acquired (yield True). This lets
+#     shared code run unchanged across backends.
+#
+# Usage (both writers must use the SAME key format):
+#   from tradedangerous.db.locks import station_advisory_lock
+#
+#   with sa_session_local(session_factory) as s:
+#       # (optional) set isolation once per process elsewhere:
+#       # s.execute(text("SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED")); s.commit()
+#       with station_advisory_lock(s, station_id, timeout_seconds=0.2, max_retries=4) as got:
+#           if not got:
+#               # processor: defer/requeue work for this station and continue
+#               return
+#           with s.begin():
+#               # do per-station writes here...
+#               pass
+# -----------------------------------------------------------------------------
+
+from __future__ import annotations
+
+import time
+from contextlib import contextmanager
+from typing import Iterator
+
+from sqlalchemy import text
+from sqlalchemy.orm import Session
+
+__all__ = [
+    "station_advisory_lock",
+    "acquire_station_lock",
+    "release_station_lock",
+    "station_lock_key",
+]
+
+# Precompiled SQL (MySQL/MariaDB only)
+_SQL_GET_LOCK     = text("SELECT GET_LOCK(:k, :t)")
+_SQL_RELEASE_LOCK = text("SELECT RELEASE_LOCK(:k)")
+
+def _is_lock_supported(session: Session) -> bool:
+    """
+    Return True if the current SQLAlchemy session is bound to a backend that
+    supports advisory locks via GET_LOCK/RELEASE_LOCK (MySQL/MariaDB).
+    """
+    try:
+        name = (session.get_bind().dialect.name or "").lower()
+    except Exception:
+        name = ""
+    return name in ("mysql", "mariadb")
+
+def _ensure_read_committed(session: Session) -> None:
+    """
+    Ensure the session is using READ COMMITTED for subsequent transactions.
+    - Applies only to MySQL/MariaDB.
+    - No-ops on SQLite/others.
+    - Only sets it if NOT already inside a transaction (affects next txn).
+    """
+    if not _is_lock_supported(session):
+        return
+    try:
+        # Only set if we're not already in a transaction; otherwise it would
+        # affect the next transaction, not the current one.
+        if not session.in_transaction():
+            session.execute(text("SET SESSION TRANSACTION ISOLATION LEVEL READ COMMITTED"))
+            # No explicit commit needed; this is a session-level setting.
+    except Exception:
+        # Best-effort; if this fails we just proceed with the default isolation.
+        pass
+
+def station_lock_key(station_id: int) -> str:
+    """
+    Return the advisory lock key used by both writers for the same station.
+    Keep this format identical in all writers (processor + Spansh).
+    """
+    return f"td.station.{int(station_id)}"
+
+def acquire_station_lock(session: Session, station_id: int, timeout_seconds: float) -> bool:
+    """
+    Try to acquire the advisory lock for a station on THIS DB connection.
+    
+    Returns:
+        True  -> acquired within timeout (or NO-OP True on unsupported dialects)
+        False -> timed out (lock held elsewhere)
+    
+    Notes:
+        - Advisory locks are per-connection. Use the same Session for acquire,
+          the critical section, and release.
+        - On SQLite/unsupported dialects, this is a NO-OP that returns True.
+    """
+    if not _is_lock_supported(session):
+        return True  # NO-OP on SQLite/unsupported backends
+    
+    key = station_lock_key(station_id)
+    row = session.execute(_SQL_GET_LOCK, {"k": key, "t": float(timeout_seconds)}).first()
+    # MariaDB/MySQL GET_LOCK returns 1 (acquired), 0 (timeout), or NULL (error)
+    return bool(row and row[0] == 1)
+
+def release_station_lock(session: Session, station_id: int) -> None:
+    """
+    Release the advisory lock for a station on THIS DB connection.
+    Safe to call in finally; releasing a non-held lock is harmless.
+    
+    On SQLite/unsupported dialects, this is a NO-OP.
+    """
+    if not _is_lock_supported(session):
+        return  # NO-OP on SQLite/unsupported backends
+    
+    key = station_lock_key(station_id)
+    try:
+        session.execute(_SQL_RELEASE_LOCK, {"k": key})
+    except Exception:
+        # Intentionally swallow — RELEASE_LOCK may return 0/NULL if not held.
+        pass
+
+@contextmanager
+def station_advisory_lock(
+    session: Session,
+    station_id: int,
+    timeout_seconds: float = 0.2,
+    max_retries: int = 4,
+    backoff_start_seconds: float = 0.05,
+) -> Iterator[bool]:
+    """
+    Context manager to acquire/retry/release a per-station advisory lock.
+    
+    Deadlock-safety requirement:
+      - Do NOT release the advisory lock before the station's writes are COMMITTED.
+      - Previously we only committed when this helper created the transaction.
+        If the Session already had an active transaction (SQLAlchemy autobegin),
+        the lock could be released while row locks were still pending commit.
+    
+    Behaviour:
+      - On MySQL/MariaDB: tries GET_LOCK() with bounded retries + exponential backoff.
+      - If acquired (got=True): COMMIT on normal exit BEFORE releasing the advisory lock,
+        regardless of whether this helper started the transaction.
+      - If NOT acquired (got=False) and this helper started the transaction: ROLLBACK to
+        avoid leaving an idle open transaction pinned to a connection.
+      - If an exception escapes the caller's block: ROLLBACK (best-effort) then re-raise.
+      - On unsupported dialects (e.g. SQLite): yields True and does nothing.
+    
+    WARNING:
+      - Do not wrap this context manager inside an external transaction manager
+        (e.g. `with session.begin():`) because it may COMMIT inside that scope.
+    """
+    # Fast-path NO-OP for SQLite/unsupported dialects
+    if not _is_lock_supported(session):
+        yield True
+        return
+    
+    # Prefer READ COMMITTED to reduce lock contention (best-effort).
+    _ensure_read_committed(session)
+    
+    started_txn = False
+    txn_ctx = None
+    if not session.in_transaction():
+        # Pin lock + DML to the same connection by opening a txn.
+        txn_ctx = session.begin()
+        started_txn = True
+    
+    got = False
+    try:
+        attempt = 0
+        while attempt < max_retries:
+            if acquire_station_lock(session, station_id, timeout_seconds):
+                got = True
+                break
+            time.sleep(backoff_start_seconds * (2 ** attempt))
+            attempt += 1
+        
+        # Hand control to caller
+        yield got
+        
+        if got:
+            # Commit while the advisory lock is still held.
+            if session.in_transaction():
+                session.commit()
+        else:
+            # If we opened a txn just to attempt locking, close it out cleanly.
+            if started_txn and session.in_transaction():
+                session.rollback()
+    
+    except Exception:
+        # Ensure we don't leak row locks / open txn on error.
+        if session.in_transaction():
+            try:
+                session.rollback()
+            except Exception:
+                pass
+        raise
+    
+    finally:
+        # Release advisory lock after commit/rollback decisions above.
+        if got:
+            try:
+                release_station_lock(session, station_id)
+            except Exception:
+                pass
+        
+        if started_txn and txn_ctx is not None:
+            try:
+                txn_ctx.close()
+            except Exception:
+                pass