<s> """<STR_LIT>""" <EOL> import json <EOL> import logging <EOL> import os <EOL> import re <EOL> import sys <EOL> import urllib2 <EOL> import xml . etree . ElementTree as et <EOL> sys . path . append ( os . path . join ( os . path . dirname ( __file__ ) , "<STR_LIT>" ) ) <EOL> from google . appengine . api import queueinfo <EOL> from google . appengine . api import datastore <EOL> from google . appengine . api import datastore_types <EOL> sys . path . append ( os . path . join ( os . path . dirname ( __file__ ) , "<STR_LIT>" ) ) <EOL> import appscale_info <EOL> import file_io <EOL> import xmltodict <EOL> class TaskQueueConfig ( ) : <EOL> """<STR_LIT>""" <EOL> QUEUE_KIND = "<STR_LIT>" <EOL> RABBITMQ = <NUM_LIT:0> <EOL> CELERY_CONCURRENCY = <NUM_LIT:10> <EOL> DEFAULT_QUEUE_YAML = """<STR_LIT>""" <EOL> DEFAULT_RATE = "<STR_LIT>" <EOL> APPSCALE_QUEUES = "<STR_LIT>" <EOL> QUEUE_INFO = "<STR_LIT>" <EOL> APP_NAME = "<STR_LIT>" <EOL> QUEUE_INFO_DB = <NUM_LIT:0> <EOL> QUEUE_INFO_FILE = <NUM_LIT:1> <EOL> CELERY_CONFIG_DIR = '<STR_LIT>' <EOL> CELERY_WORKER_DIR = '<STR_LIT>' <EOL> CELERY_STATE_DIR = '<STR_LIT>' <EOL> TEMPLATE_DIR = os . path . dirname ( os . path . realpath ( __file__ ) ) + "<STR_LIT>" <EOL> HEADER_LOC = TEMPLATE_DIR + '<STR_LIT>' <EOL> TASK_LOC = TEMPLATE_DIR + '<STR_LIT>' <EOL> MAX_QUEUE_NAME_LENGTH = <NUM_LIT:100> <EOL> QUEUE_NAME_PATTERN = r'<STR_LIT>' % MAX_QUEUE_NAME_LENGTH <EOL> QUEUE_NAME_RE = re . compile ( QUEUE_NAME_PATTERN ) <EOL> YAML_TO_XML_TAGS_TO_CONVERT = [ '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' ] <EOL> TAGS_TO_IGNORE = [ '<STR_LIT>' ] <EOL> def __init__ ( self , broker , app_id ) : <EOL> """<STR_LIT>""" <EOL> file_io . set_logging_format ( ) <EOL> self . _broker = broker <EOL> self . _broker_location = self . __broker_location ( broker ) <EOL> self . _app_id = app_id <EOL> self . _queue_info_db = None <EOL> self . _queue_info_file = None <EOL> file_io . mkdir ( self . CELERY_CONFIG_DIR ) <EOL> file_io . mkdir ( self . CELERY_WORKER_DIR ) <EOL> def get_queue_file_location ( self , app_id ) : <EOL> """<STR_LIT>""" <EOL> queue_yaml = appscale_info . get_app_path ( app_id ) + '<STR_LIT>' <EOL> queue_xml = appscale_info . get_app_path ( app_id ) + '<STR_LIT>' <EOL> if file_io . exists ( queue_yaml ) : <EOL> return queue_yaml <EOL> elif file_io . exists ( queue_xml ) : <EOL> return queue_xml <EOL> else : <EOL> return "<STR_LIT>" <EOL> def load_queues_from_file ( self , app_id ) : <EOL> """<STR_LIT>""" <EOL> queue_file = self . get_queue_file_location ( app_id ) <EOL> info = "<STR_LIT>" <EOL> using_default = False <EOL> try : <EOL> info = file_io . read ( queue_file ) <EOL> logging . info ( "<STR_LIT>" . format ( app_id ) ) <EOL> except IOError : <EOL> logging . info ( "<STR_LIT>" . format ( app_id ) ) <EOL> info = self . DEFAULT_QUEUE_YAML <EOL> using_default = True <EOL> queue_info = "<STR_LIT>" <EOL> if queue_file . endswith ( '<STR_LIT>' ) or using_default : <EOL> queue_info = queueinfo . LoadSingleQueue ( info ) . ToDict ( ) <EOL> elif queue_file . endswith ( '<STR_LIT>' ) : <EOL> queue_info = self . parse_queue_xml ( info ) <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % queue_file ) <EOL> if not queue_info : <EOL> raise ValueError ( "<STR_LIT>" % queue_file ) <EOL> has_default = False <EOL> if '<STR_LIT>' not in queue_info or len ( queue_info [ '<STR_LIT>' ] ) == <NUM_LIT:0> : <EOL> queue_info = { '<STR_LIT>' : [ { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:name>' : '<STR_LIT:default>' } ] } <EOL> for queue in queue_info [ '<STR_LIT>' ] : <EOL> if queue [ '<STR_LIT:name>' ] == '<STR_LIT:default>' : <EOL> has_default = True <EOL> if not has_default : <EOL> queue_info [ '<STR_LIT>' ] . append ( { '<STR_LIT>' : '<STR_LIT>' , '<STR_LIT:name>' : '<STR_LIT:default>' } ) <EOL> self . _queue_info_file = queue_info <EOL> logging . info ( "<STR_LIT>" . format ( app_id , queue_info ) ) <EOL> return queue_info <EOL> def parse_queue_xml ( self , xml_string ) : <EOL> """<STR_LIT>""" <EOL> xml_dict = xmltodict . parse ( xml_string ) <EOL> converted = { '<STR_LIT>' : [ ] } <EOL> if isinstance ( xml_dict [ '<STR_LIT>' ] [ '<STR_LIT>' ] , list ) : <EOL> all_queues = xml_dict [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> else : <EOL> all_queues = [ xml_dict [ '<STR_LIT>' ] [ '<STR_LIT>' ] ] <EOL> for queue in all_queues : <EOL> single_queue = { } <EOL> for tag , value in queue . iteritems ( ) : <EOL> if tag in self . YAML_TO_XML_TAGS_TO_CONVERT : <EOL> tag = tag . replace ( '<STR_LIT:->' , '<STR_LIT:_>' ) <EOL> if tag == "<STR_LIT>" : <EOL> retry_dict = { } <EOL> for retry_tag in queue [ '<STR_LIT>' ] : <EOL> value = queue [ '<STR_LIT>' ] [ retry_tag ] <EOL> if retry_tag in self . YAML_TO_XML_TAGS_TO_CONVERT : <EOL> retry_tag = retry_tag . replace ( '<STR_LIT:->' , '<STR_LIT:_>' ) <EOL> if retry_tag not in self . TAGS_TO_IGNORE : <EOL> retry_dict [ str ( retry_tag ) ] = str ( value ) . strip ( '<STR_LIT>' ) <EOL> single_queue [ '<STR_LIT>' ] = retry_dict <EOL> elif tag not in self . TAGS_TO_IGNORE : <EOL> single_queue [ str ( tag ) ] = str ( value ) . strip ( '<STR_LIT>' ) <EOL> converted [ '<STR_LIT>' ] . append ( single_queue ) <EOL> logging . debug ( "<STR_LIT>" . format ( converted ) ) <EOL> return converted <EOL> def get_file_queue_info ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _queue_info_file <EOL> def get_db_queue_info ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _queue_info_db <EOL> def __get_queues_from_db ( self ) : <EOL> """<STR_LIT>""" <EOL> queues_key = datastore . Key . from_path ( self . QUEUE_KIND , <EOL> self . _app_id , <EOL> _app = self . APPSCALE_QUEUES ) <EOL> queues = datastore . Get ( queues_key ) <EOL> return json . loads ( queues [ self . QUEUE_INFO ] ) <EOL> def load_queues_from_db ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _queue_info_db = self . __get_queues_from_db ( ) <EOL> return self . _queue_info_db <EOL> def save_queues_to_db ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . _queue_info_file : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> json_queues = json . dumps ( self . _queue_info_file ) <EOL> entity = datastore . Entity ( self . QUEUE_KIND , <EOL> name = self . _app_id , <EOL> _app = self . APPSCALE_QUEUES ) <EOL> entity . update ( { self . QUEUE_INFO : datastore_types . Blob ( json_queues ) , <EOL> self . APP_NAME : datastore_types . ByteString ( self . _app_id ) } ) <EOL> datastore . Put ( entity ) <EOL> def create_celery_worker_scripts ( self , input_type ) : <EOL> """<STR_LIT>""" <EOL> queue_info = self . _queue_info_file <EOL> if input_type == self . QUEUE_INFO_DB : <EOL> queue_info = self . _queue_info_db <EOL> header_template = file_io . read ( self . HEADER_LOC ) <EOL> task_template = file_io . read ( self . TASK_LOC ) <EOL> header_template = header_template . replace ( "<STR_LIT>" , self . _app_id ) <EOL> script = header_template . replace ( "<STR_LIT>" , <EOL> self . _app_id ) + '<STR_LIT:\n>' <EOL> for queue in queue_info [ '<STR_LIT>' ] : <EOL> queue_name = queue [ '<STR_LIT:name>' ] <EOL> queue_name = queue_name . replace ( '<STR_LIT:->' , '<STR_LIT:_>' ) <EOL> self . validate_queue_name ( queue_name ) <EOL> new_task = task_template . replace ( "<STR_LIT>" , <EOL> self . get_queue_function_name ( queue_name ) ) <EOL> script += new_task + '<STR_LIT:\n>' <EOL> worker_file = self . get_celery_worker_script_path ( self . _app_id ) <EOL> file_io . write ( worker_file , script ) <EOL> return worker_file <EOL> @ staticmethod <EOL> def remove_config_files ( app_id ) : <EOL> """<STR_LIT>""" <EOL> worker_file = TaskQueueConfig . get_celery_worker_script_path ( app_id ) <EOL> config_file = TaskQueueConfig . get_celery_configuration_path ( app_id ) <EOL> file_io . delete ( worker_file ) <EOL> file_io . delete ( config_file ) <EOL> @ staticmethod <EOL> def get_queue_function_name ( queue_name ) : <EOL> """<STR_LIT>""" <EOL> queue_name = queue_name . replace ( '<STR_LIT:->' , '<STR_LIT:_>' ) <EOL> return "<STR_LIT>" % queue_name <EOL> @ staticmethod <EOL> def get_celery_annotation_name ( app_id , queue_name ) : <EOL> """<STR_LIT>""" <EOL> module_name = TaskQueueConfig . get_celery_worker_module_name ( app_id ) <EOL> function_name = TaskQueueConfig . get_queue_function_name ( queue_name ) <EOL> return "<STR_LIT>" % ( module_name , function_name ) <EOL> @ staticmethod <EOL> def get_celery_worker_script_path ( app_id ) : <EOL> """<STR_LIT>""" <EOL> return TaskQueueConfig . CELERY_WORKER_DIR + "<STR_LIT>" + app_id + "<STR_LIT>" <EOL> @ staticmethod <EOL> def get_celery_worker_module_name ( app_id ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT>" + app_id <EOL> @ staticmethod <EOL> def get_celery_configuration_path ( app_id ) : <EOL> """<STR_LIT>""" <EOL> return TaskQueueConfig . CELERY_CONFIG_DIR + app_id + "<STR_LIT>" <EOL> def create_celery_file ( self , input_type ) : <EOL> """<STR_LIT>""" <EOL> queue_info = self . _queue_info_file <EOL> if input_type == self . QUEUE_INFO_DB : <EOL> queue_info = self . _queue_info_db <EOL> celery_queues = [ ] <EOL> celery_annotations = [ ] <EOL> for queue in queue_info [ '<STR_LIT>' ] : <EOL> if '<STR_LIT>' in queue and queue [ '<STR_LIT>' ] == "<STR_LIT>" : <EOL> continue <EOL> celery_queue_name = TaskQueueConfig . get_celery_queue_name ( self . _app_id , queue [ '<STR_LIT:name>' ] ) <EOL> celery_queues . append ( "<STR_LIT>" + celery_queue_name + "<STR_LIT>" + self . _app_id + "<STR_LIT>" + celery_queue_name + "<STR_LIT>" ) <EOL> rate_limit = self . DEFAULT_RATE <EOL> if '<STR_LIT>' in queue : <EOL> rate_limit = queue [ '<STR_LIT>' ] <EOL> annotation_name = TaskQueueConfig . get_celery_annotation_name ( self . _app_id , <EOL> queue [ '<STR_LIT:name>' ] ) <EOL> celery_annotations . append ( "<STR_LIT:'>" + annotation_name + "<STR_LIT>" + rate_limit + "<STR_LIT>" ) <EOL> celery_queues = '<STR_LIT:\n>' . join ( celery_queues ) <EOL> celery_annotations = '<STR_LIT:\n>' . join ( celery_annotations ) <EOL> config = """<STR_LIT>""" <EOL> config += celery_queues <EOL> config += """<STR_LIT>""" <EOL> config += celery_annotations <EOL> config += """<STR_LIT>""" <EOL> config_file = self . _app_id + "<STR_LIT>" <EOL> file_io . write ( self . CELERY_CONFIG_DIR + config_file , config ) <EOL> return self . CELERY_CONFIG_DIR + config_file <EOL> def __broker_location ( self , broker ) : <EOL> """<STR_LIT>""" <EOL> if broker == self . RABBITMQ : <EOL> from brokers import rabbitmq <EOL> return rabbitmq . get_connection_string ( ) <EOL> else : <EOL> raise NotImplementedError ( <EOL> "<STR_LIT>" % broker ) <EOL> def get_broker_string ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _broker_location <EOL> def validate_queue_name ( self , queue_name ) : <EOL> """<STR_LIT>""" <EOL> if not self . QUEUE_NAME_RE . match ( queue_name ) : <EOL> raise NameError ( "<STR_LIT>" % ( queue_name , self . QUEUE_NAME_PATTERN ) ) <EOL> @ staticmethod <EOL> def get_celery_queue_name ( app_id , queue_name ) : <EOL> """<STR_LIT>""" <EOL> return app_id + "<STR_LIT>" + queue_name </s>
