<s> """<STR_LIT>""" <EOL> import datetime <EOL> import json <EOL> import logging <EOL> import time <EOL> import uuid <EOL> from . import app <EOL> from . import db <EOL> from dpxdt . server import signals <EOL> class Error ( Exception ) : <EOL> """<STR_LIT>""" <EOL> class TaskDoesNotExistError ( Error ) : <EOL> """<STR_LIT>""" <EOL> class LeaseExpiredError ( Error ) : <EOL> """<STR_LIT>""" <EOL> class NotOwnerError ( Error ) : <EOL> """<STR_LIT>""" <EOL> class WorkQueue ( db . Model ) : <EOL> """<STR_LIT>""" <EOL> CANCELED = '<STR_LIT>' <EOL> DONE = '<STR_LIT>' <EOL> ERROR = '<STR_LIT:error>' <EOL> LIVE = '<STR_LIT>' <EOL> STATES = frozenset ( [ CANCELED , DONE , ERROR , LIVE ] ) <EOL> task_id = db . Column ( db . String ( <NUM_LIT:100> ) , primary_key = True , nullable = False ) <EOL> queue_name = db . Column ( db . String ( <NUM_LIT:100> ) , primary_key = True , nullable = False ) <EOL> status = db . Column ( db . Enum ( * STATES ) , default = LIVE , nullable = False ) <EOL> eta = db . Column ( db . DateTime , default = datetime . datetime . utcnow , <EOL> nullable = False ) <EOL> build_id = db . Column ( db . Integer , db . ForeignKey ( '<STR_LIT>' ) ) <EOL> release_id = db . Column ( db . Integer , db . ForeignKey ( '<STR_LIT>' ) ) <EOL> run_id = db . Column ( db . Integer , db . ForeignKey ( '<STR_LIT>' ) ) <EOL> source = db . Column ( db . String ( <NUM_LIT> ) ) <EOL> created = db . Column ( db . DateTime , default = datetime . datetime . utcnow ) <EOL> finished = db . Column ( db . DateTime ) <EOL> lease_attempts = db . Column ( db . Integer , default = <NUM_LIT:0> , nullable = False ) <EOL> last_lease = db . Column ( db . DateTime ) <EOL> last_owner = db . Column ( db . String ( <NUM_LIT> ) ) <EOL> heartbeat = db . Column ( db . Text ) <EOL> heartbeat_number = db . Column ( db . Integer ) <EOL> payload = db . Column ( db . LargeBinary ) <EOL> content_type = db . Column ( db . String ( <NUM_LIT:100> ) ) <EOL> __table_args__ = ( <EOL> db . Index ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:status>' , '<STR_LIT>' ) , <EOL> db . Index ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT:status>' , '<STR_LIT>' ) , <EOL> db . Index ( '<STR_LIT>' , '<STR_LIT:status>' , '<STR_LIT>' ) , <EOL> ) <EOL> @ property <EOL> def lease_outstanding ( self ) : <EOL> if not self . status == WorkQueue . LIVE : <EOL> return False <EOL> if not self . last_owner : <EOL> return False <EOL> now = datetime . datetime . utcnow ( ) <EOL> return now < self . eta <EOL> def add ( queue_name , payload = None , content_type = None , source = None , task_id = None , <EOL> build_id = None , release_id = None , run_id = None ) : <EOL> """<STR_LIT>""" <EOL> if task_id : <EOL> task = WorkQueue . query . filter_by ( task_id = task_id ) . first ( ) <EOL> if task : <EOL> return task . task_id <EOL> else : <EOL> task_id = uuid . uuid4 ( ) . hex <EOL> if payload and not content_type and not isinstance ( payload , basestring ) : <EOL> payload = json . dumps ( payload ) <EOL> content_type = '<STR_LIT:application/json>' <EOL> now = datetime . datetime . utcnow ( ) <EOL> task = WorkQueue ( <EOL> task_id = task_id , <EOL> queue_name = queue_name , <EOL> eta = now , <EOL> source = source , <EOL> build_id = build_id , <EOL> release_id = release_id , <EOL> run_id = run_id , <EOL> payload = payload , <EOL> content_type = content_type ) <EOL> db . session . add ( task ) <EOL> return task . task_id <EOL> def _datetime_to_epoch_seconds ( dt ) : <EOL> """<STR_LIT>""" <EOL> if dt is None : <EOL> return None <EOL> return int ( time . mktime ( dt . utctimetuple ( ) ) ) <EOL> def _task_to_dict ( task ) : <EOL> """<STR_LIT>""" <EOL> payload = task . payload <EOL> if payload and task . content_type == '<STR_LIT:application/json>' : <EOL> payload = json . loads ( payload ) <EOL> return dict ( <EOL> task_id = task . task_id , <EOL> queue_name = task . queue_name , <EOL> eta = _datetime_to_epoch_seconds ( task . eta ) , <EOL> source = task . source , <EOL> created = _datetime_to_epoch_seconds ( task . created ) , <EOL> lease_attempts = task . lease_attempts , <EOL> last_lease = _datetime_to_epoch_seconds ( task . last_lease ) , <EOL> payload = payload , <EOL> content_type = task . content_type ) <EOL> def lease ( queue_name , owner , count = <NUM_LIT:1> , timeout_seconds = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> now = datetime . datetime . utcnow ( ) <EOL> query = ( <EOL> WorkQueue . query <EOL> . filter_by ( queue_name = queue_name , status = WorkQueue . LIVE ) <EOL> . filter ( WorkQueue . eta <= now ) <EOL> . order_by ( WorkQueue . eta ) <EOL> . with_lockmode ( '<STR_LIT>' ) <EOL> . limit ( count ) ) <EOL> task_list = query . all ( ) <EOL> if not task_list : <EOL> return None <EOL> next_eta = now + datetime . timedelta ( seconds = timeout_seconds ) <EOL> for task in task_list : <EOL> task . eta = next_eta <EOL> task . lease_attempts += <NUM_LIT:1> <EOL> task . last_owner = owner <EOL> task . last_lease = now <EOL> task . heartbeat = None <EOL> task . heartbeat_number = <NUM_LIT:0> <EOL> db . session . add ( task ) <EOL> return [ _task_to_dict ( task ) for task in task_list ] <EOL> def _get_task_with_policy ( queue_name , task_id , owner ) : <EOL> """<STR_LIT>""" <EOL> now = datetime . datetime . utcnow ( ) <EOL> task = ( <EOL> WorkQueue . query <EOL> . filter_by ( queue_name = queue_name , task_id = task_id ) <EOL> . with_lockmode ( '<STR_LIT>' ) <EOL> . first ( ) ) <EOL> if not task : <EOL> raise TaskDoesNotExistError ( '<STR_LIT>' % task_id ) <EOL> lease_delta = now - task . eta <EOL> if lease_delta > datetime . timedelta ( <NUM_LIT:0> ) : <EOL> db . session . rollback ( ) <EOL> raise LeaseExpiredError ( '<STR_LIT>' % ( <EOL> task . queue_name , task_id , lease_delta ) ) <EOL> if task . last_owner != owner : <EOL> db . session . rollback ( ) <EOL> raise NotOwnerError ( '<STR_LIT>' % ( <EOL> task . queue_name , task_id , task . last_owner ) ) <EOL> return task <EOL> def heartbeat ( queue_name , task_id , owner , message , index ) : <EOL> """<STR_LIT>""" <EOL> task = _get_task_with_policy ( queue_name , task_id , owner ) <EOL> if task . heartbeat_number > index : <EOL> return False <EOL> task . heartbeat = message <EOL> task . heartbeat_number = index <EOL> now = datetime . datetime . utcnow ( ) <EOL> timeout_delta = task . eta - task . last_lease <EOL> task . eta = now + timeout_delta <EOL> task . last_lease = now <EOL> db . session . add ( task ) <EOL> signals . task_updated . send ( app , task = task ) <EOL> return True <EOL> def finish ( queue_name , task_id , owner , error = False ) : <EOL> """<STR_LIT>""" <EOL> task = _get_task_with_policy ( queue_name , task_id , owner ) <EOL> if not task . status == WorkQueue . LIVE : <EOL> logging . warning ( '<STR_LIT>' <EOL> '<STR_LIT>' , <EOL> task . queue_name , task_id , owner , task . status ) <EOL> return False <EOL> if not error : <EOL> task . status = WorkQueue . DONE <EOL> else : <EOL> task . status = WorkQueue . ERROR <EOL> task . finished = datetime . datetime . utcnow ( ) <EOL> db . session . add ( task ) <EOL> signals . task_updated . send ( app , task = task ) <EOL> return True <EOL> def _query ( queue_name = None , build_id = None , release_id = None , run_id = None , <EOL> count = None ) : <EOL> """<STR_LIT>""" <EOL> assert queue_name or build_id or release_id or run_id <EOL> q = WorkQueue . query <EOL> if queue_name : <EOL> q = q . filter_by ( queue_name = queue_name ) <EOL> if build_id : <EOL> q = q . filter_by ( build_id = build_id ) <EOL> if release_id : <EOL> q = q . filter_by ( release_id = release_id ) <EOL> if run_id : <EOL> q = q . filter_by ( run_id = run_id ) <EOL> q = q . order_by ( WorkQueue . created . desc ( ) ) <EOL> if count is not None : <EOL> q = q . limit ( count ) <EOL> return q . all ( ) <EOL> def query ( ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> count = kwargs . get ( '<STR_LIT:count>' , None ) <EOL> task_list = _query ( ** kwargs ) <EOL> task_dict_list = [ _task_to_dict ( task ) for task in task_list ] <EOL> if count == <NUM_LIT:1> : <EOL> if not task_dict_list : <EOL> return None <EOL> else : <EOL> return task_dict_list [ <NUM_LIT:0> ] <EOL> return task_dict_list <EOL> def cancel ( ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> task_list = _query ( ** kwargs ) <EOL> for task in task_list : <EOL> task . status = WorkQueue . CANCELED <EOL> task . finished = datetime . datetime . utcnow ( ) <EOL> db . session . add ( task ) <EOL> return len ( task_list ) </s>
