101 lines
3.3 KiB
Python
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()
|