<s> import os <EOL> import sys <EOL> import signal <EOL> import ConfigParser <EOL> import redis <EOL> from sharq . utils import ( is_valid_identifier , is_valid_interval , <EOL> is_valid_requeue_limit , generate_epoch , <EOL> serialize_payload , deserialize_payload ) <EOL> from sharq . exceptions import SharqException , BadArgumentException <EOL> class SharQ ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , config_path ) : <EOL> """<STR_LIT>""" <EOL> self . config_path = config_path <EOL> self . _load_config ( ) <EOL> self . _initialize ( ) <EOL> def _initialize ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _key_prefix = self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> self . _job_expire_interval = int ( <EOL> self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> ) <EOL> self . _default_job_requeue_limit = int ( <EOL> self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> ) <EOL> redis_connection_type = self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> db = self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if redis_connection_type == '<STR_LIT>' : <EOL> self . _r = redis . StrictRedis ( <EOL> db = db , <EOL> unix_socket_path = self . _config . get ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> ) <EOL> elif redis_connection_type == '<STR_LIT>' : <EOL> self . _r = redis . StrictRedis ( <EOL> db = db , <EOL> host = self . _config . get ( '<STR_LIT>' , '<STR_LIT:host>' ) , <EOL> port = self . _config . get ( '<STR_LIT>' , '<STR_LIT:port>' ) <EOL> ) <EOL> self . _load_lua_scripts ( ) <EOL> def _load_config ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _config = ConfigParser . SafeConfigParser ( ) <EOL> self . _config . read ( self . config_path ) <EOL> def reload_config ( self , config_path = None ) : <EOL> """<STR_LIT>""" <EOL> if config_path : <EOL> self . config_path = config_path <EOL> self . _load_config ( ) <EOL> def _load_lua_scripts ( self ) : <EOL> """<STR_LIT>""" <EOL> lua_script_path = os . path . join ( <EOL> os . path . dirname ( os . path . abspath ( __file__ ) ) , <EOL> '<STR_LIT>' <EOL> ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as enqueue_file : <EOL> self . _lua_enqueue_script = enqueue_file . read ( ) <EOL> self . _lua_enqueue = self . _r . register_script ( <EOL> self . _lua_enqueue_script ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as dequeue_file : <EOL> self . _lua_dequeue_script = dequeue_file . read ( ) <EOL> self . _lua_dequeue = self . _r . register_script ( <EOL> self . _lua_dequeue_script ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as finish_file : <EOL> self . _lua_finish_script = finish_file . read ( ) <EOL> self . _lua_finish = self . _r . register_script ( self . _lua_finish_script ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as interval_file : <EOL> self . _lua_interval_script = interval_file . read ( ) <EOL> self . _lua_interval = self . _r . register_script ( <EOL> self . _lua_interval_script ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as requeue_file : <EOL> self . _lua_requeue_script = requeue_file . read ( ) <EOL> self . _lua_requeue = self . _r . register_script ( <EOL> self . _lua_requeue_script ) <EOL> with open ( os . path . join ( <EOL> lua_script_path , <EOL> '<STR_LIT>' ) , '<STR_LIT:r>' ) as metrics_file : <EOL> self . _lua_metrics_script = metrics_file . read ( ) <EOL> self . _lua_metrics = self . _r . register_script ( <EOL> self . _lua_metrics_script ) <EOL> def reload_lua_scripts ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _load_lua_scripts ( ) <EOL> def enqueue ( self , payload , interval , job_id , <EOL> queue_id , queue_type = '<STR_LIT:default>' , requeue_limit = None ) : <EOL> """<STR_LIT>""" <EOL> if not is_valid_interval ( interval ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( job_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_type ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if requeue_limit is None : <EOL> requeue_limit = self . _default_job_requeue_limit <EOL> if not is_valid_requeue_limit ( requeue_limit ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> try : <EOL> serialized_payload = serialize_payload ( payload ) <EOL> except TypeError as e : <EOL> raise BadArgumentException ( e . message ) <EOL> timestamp = str ( generate_epoch ( ) ) <EOL> keys = [ <EOL> self . _key_prefix , <EOL> queue_type <EOL> ] <EOL> args = [ <EOL> timestamp , <EOL> queue_id , <EOL> job_id , <EOL> '<STR_LIT>' % serialized_payload , <EOL> interval , <EOL> requeue_limit <EOL> ] <EOL> self . _lua_enqueue ( keys = keys , args = args ) <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT>' <EOL> } <EOL> return response <EOL> def dequeue ( self , queue_type = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> if not is_valid_identifier ( queue_type ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> timestamp = str ( generate_epoch ( ) ) <EOL> keys = [ <EOL> self . _key_prefix , <EOL> queue_type <EOL> ] <EOL> args = [ <EOL> timestamp , <EOL> self . _job_expire_interval <EOL> ] <EOL> dequeue_response = self . _lua_dequeue ( keys = keys , args = args ) <EOL> if len ( dequeue_response ) < <NUM_LIT:4> : <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT>' <EOL> } <EOL> return response <EOL> queue_id , job_id , payload , requeues_remaining = dequeue_response <EOL> payload = deserialize_payload ( payload [ <NUM_LIT:1> : - <NUM_LIT:1> ] ) <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' , <EOL> '<STR_LIT>' : queue_id , <EOL> '<STR_LIT>' : job_id , <EOL> '<STR_LIT>' : payload , <EOL> '<STR_LIT>' : int ( requeues_remaining ) <EOL> } <EOL> return response <EOL> def finish ( self , job_id , queue_id , queue_type = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> if not is_valid_identifier ( job_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_type ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> keys = [ <EOL> self . _key_prefix , <EOL> queue_type <EOL> ] <EOL> args = [ <EOL> queue_id , <EOL> job_id <EOL> ] <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' <EOL> } <EOL> finish_response = self . _lua_finish ( keys = keys , args = args ) <EOL> if finish_response == <NUM_LIT:0> : <EOL> response . update ( { <EOL> '<STR_LIT:status>' : '<STR_LIT>' <EOL> } ) <EOL> return response <EOL> def interval ( self , interval , queue_id , queue_type = '<STR_LIT:default>' ) : <EOL> """<STR_LIT>""" <EOL> if not is_valid_interval ( interval ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if not is_valid_identifier ( queue_type ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> interval_hmap_key = '<STR_LIT>' % self . _key_prefix <EOL> interval_queue_key = '<STR_LIT>' % ( queue_type , queue_id ) <EOL> keys = [ <EOL> interval_hmap_key , <EOL> interval_queue_key <EOL> ] <EOL> args = [ <EOL> interval <EOL> ] <EOL> interval_response = self . _lua_interval ( keys = keys , args = args ) <EOL> if interval_response == <NUM_LIT:0> : <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT>' <EOL> } <EOL> else : <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' <EOL> } <EOL> return response <EOL> def requeue ( self ) : <EOL> """<STR_LIT>""" <EOL> timestamp = str ( generate_epoch ( ) ) <EOL> active_queue_type_list = self . _r . smembers ( <EOL> '<STR_LIT>' % self . _key_prefix ) <EOL> for queue_type in active_queue_type_list : <EOL> keys = [ <EOL> self . _key_prefix , <EOL> queue_type <EOL> ] <EOL> args = [ <EOL> timestamp <EOL> ] <EOL> job_discard_list = self . _lua_requeue ( keys = keys , args = args ) <EOL> for job in job_discard_list : <EOL> queue_id , job_id = job . split ( '<STR_LIT::>' ) <EOL> self . finish ( <EOL> job_id = job_id , <EOL> queue_id = queue_id , <EOL> queue_type = queue_type <EOL> ) <EOL> def metrics ( self , queue_type = None , queue_id = None ) : <EOL> """<STR_LIT>""" <EOL> if queue_id is not None and not is_valid_identifier ( queue_id ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> if queue_type is not None and not is_valid_identifier ( queue_type ) : <EOL> raise BadArgumentException ( '<STR_LIT>' ) <EOL> response = { <EOL> '<STR_LIT:status>' : '<STR_LIT>' <EOL> } <EOL> if not queue_type and not queue_id : <EOL> active_queue_types = self . _r . smembers ( <EOL> '<STR_LIT>' % self . _key_prefix ) <EOL> ready_queue_types = self . _r . smembers ( <EOL> '<STR_LIT>' % self . _key_prefix ) <EOL> all_queue_types = active_queue_types | ready_queue_types <EOL> timestamp = str ( generate_epoch ( ) ) <EOL> keys = [ <EOL> self . _key_prefix <EOL> ] <EOL> args = [ <EOL> timestamp <EOL> ] <EOL> enqueue_details , dequeue_details = self . _lua_metrics ( <EOL> keys = keys , args = args ) <EOL> enqueue_counts = { } <EOL> dequeue_counts = { } <EOL> for i in xrange ( <NUM_LIT:0> , len ( enqueue_details ) , <NUM_LIT:2> ) : <EOL> enqueue_counts [ str ( enqueue_details [ i ] ) ] = int ( <EOL> enqueue_details [ i + <NUM_LIT:1> ] or <NUM_LIT:0> ) <EOL> dequeue_counts [ str ( dequeue_details [ i ] ) ] = int ( <EOL> dequeue_details [ i + <NUM_LIT:1> ] or <NUM_LIT:0> ) <EOL> response . update ( { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' , <EOL> '<STR_LIT>' : list ( all_queue_types ) , <EOL> '<STR_LIT>' : enqueue_counts , <EOL> '<STR_LIT>' : dequeue_counts <EOL> } ) <EOL> return response <EOL> elif queue_type and not queue_id : <EOL> pipe = self . _r . pipeline ( ) <EOL> pipe . zrange ( '<STR_LIT>' % ( self . _key_prefix , queue_type ) , <NUM_LIT:0> , - <NUM_LIT:1> ) <EOL> pipe . zrange ( '<STR_LIT>' % ( self . _key_prefix , queue_type ) , <NUM_LIT:0> , - <NUM_LIT:1> ) <EOL> ready_queues , active_queues = pipe . execute ( ) <EOL> active_queues = [ i . split ( '<STR_LIT::>' ) [ <NUM_LIT:0> ] for i in active_queues ] <EOL> all_queue_set = set ( ready_queues ) | set ( active_queues ) <EOL> response . update ( { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' , <EOL> '<STR_LIT>' : list ( all_queue_set ) <EOL> } ) <EOL> return response <EOL> elif queue_type and queue_id : <EOL> active_queue_types = self . _r . smembers ( <EOL> '<STR_LIT>' % self . _key_prefix ) <EOL> ready_queue_types = self . _r . smembers ( <EOL> '<STR_LIT>' % self . _key_prefix ) <EOL> all_queue_types = active_queue_types | ready_queue_types <EOL> timestamp = str ( generate_epoch ( ) ) <EOL> keys = [ <EOL> '<STR_LIT>' % ( self . _key_prefix , queue_type , queue_id ) <EOL> ] <EOL> args = [ <EOL> timestamp <EOL> ] <EOL> enqueue_details , dequeue_details = self . _lua_metrics ( <EOL> keys = keys , args = args ) <EOL> enqueue_counts = { } <EOL> dequeue_counts = { } <EOL> for i in xrange ( <NUM_LIT:0> , len ( enqueue_details ) , <NUM_LIT:2> ) : <EOL> enqueue_counts [ str ( enqueue_details [ i ] ) ] = int ( <EOL> enqueue_details [ i + <NUM_LIT:1> ] or <NUM_LIT:0> ) <EOL> dequeue_counts [ str ( dequeue_details [ i ] ) ] = int ( <EOL> dequeue_details [ i + <NUM_LIT:1> ] or <NUM_LIT:0> ) <EOL> queue_length = self . _r . llen ( '<STR_LIT>' % ( <EOL> self . _key_prefix , queue_type , queue_id ) ) <EOL> response . update ( { <EOL> '<STR_LIT:status>' : '<STR_LIT:success>' , <EOL> '<STR_LIT>' : int ( queue_length ) , <EOL> '<STR_LIT>' : enqueue_counts , <EOL> '<STR_LIT>' : dequeue_counts <EOL> } ) <EOL> return response <EOL> elif not queue_type and queue_id : <EOL> raise BadArgumentException ( <EOL> '<STR_LIT>' ) <EOL> return response </s>
