atst/atat/models/utils.py
2020-03-04 11:51:15 -05:00

101 lines
3.3 KiB
Python

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()