Update atst to atat

atat/models/utils.py

@@ -0,0 +1,100 @@
+from typing import List
+
+from sqlalchemy import func, sql, Interval, and_, or_
+from contextlib import contextmanager
+
+from atat.database import db
+from atat.domain.exceptions import ClaimFailedException
+
+
+@contextmanager
+def claim_for_update(resource, minutes=30):
+    """
+    Claim a mutually exclusive expiring hold on a resource.
+    Uses the database as a central source of time in case the server clocks have drifted.
+
+    Args:
+        resource:   A SQLAlchemy model instance with a `claimed_until` attribute.
+        minutes:    The maximum amount of time, in minutes, to hold the claim.
+    """
+    Model = resource.__class__
+
+    claim_until = func.now() + func.cast(
+        sql.functions.concat(minutes, " MINUTES"), Interval
+    )
+
+    # Optimistically query for and update the resource in question. If it's
+    # already claimed, `rows_updated` will be 0 and we can give up.
+    rows_updated = (
+        db.session.query(Model)
+        .filter(
+            and_(
+                Model.id == resource.id,
+                or_(Model.claimed_until.is_(None), Model.claimed_until <= func.now()),
+            )
+        )
+        .update({"claimed_until": claim_until}, synchronize_session="fetch")
+    )
+    if rows_updated < 1:
+        raise ClaimFailedException(resource)
+
+    # Fetch the claimed resource
+    claimed = db.session.query(Model).filter_by(id=resource.id).one()
+
+    try:
+        # Give the resource to the caller.
+        yield claimed
+    finally:
+        # Release the claim.
+        db.session.query(Model).filter(Model.id == resource.id).filter(
+            Model.claimed_until != None
+        ).update({"claimed_until": None}, synchronize_session="fetch")
+        db.session.commit()
+
+
+@contextmanager
+def claim_many_for_update(resources: List, minutes=30):
+    """
+    Claim a mutually exclusive expiring hold on a group of resources.
+    Uses the database as a central source of time in case the server clocks have drifted.
+
+    Args:
+        resources:   A list of SQLAlchemy model instances with a `claimed_until` attribute.
+        minutes:    The maximum amount of time, in minutes, to hold the claim.
+    """
+    Model = resources[0].__class__
+
+    claim_until = func.now() + func.cast(
+        sql.functions.concat(minutes, " MINUTES"), Interval
+    )
+
+    ids = tuple(r.id for r in resources)
+
+    # Optimistically query for and update the resources in question. If they're
+    # already claimed, `rows_updated` will be 0 and we can give up.
+    rows_updated = (
+        db.session.query(Model)
+        .filter(
+            and_(
+                Model.id.in_(ids),
+                or_(Model.claimed_until.is_(None), Model.claimed_until <= func.now()),
+            )
+        )
+        .update({"claimed_until": claim_until}, synchronize_session="fetch")
+    )
+    if rows_updated < 1:
+        # TODO: Generalize this exception class so it can take multiple resources
+        raise ClaimFailedException(resources[0])
+
+    # Fetch the claimed resources
+    claimed = db.session.query(Model).filter(Model.id.in_(ids)).all()
+
+    try:
+        # Give the resource to the caller.
+        yield claimed
+    finally:
+        # Release the claim.
+        db.session.query(Model).filter(Model.id.in_(ids)).filter(
+            Model.claimed_until != None
+        ).update({"claimed_until": None}, synchronize_session="fetch")
+        db.session.commit()