<s> import itertools <EOL> import time <EOL> try : <EOL> import json <EOL> except ImportError : <EOL> import simplejson as json <EOL> import requests <EOL> import requests . exceptions <EOL> import marathon <EOL> from . models import MarathonApp , MarathonDeployment , MarathonGroup , MarathonInfo , MarathonTask , MarathonEndpoint , MarathonQueueItem <EOL> from . exceptions import InternalServerError , NotFoundError , MarathonHttpError , MarathonError <EOL> from . models . events import EventFactory <EOL> class MarathonClient ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , servers , username = None , password = None , timeout = <NUM_LIT:10> ) : <EOL> """<STR_LIT>""" <EOL> self . servers = servers if isinstance ( servers , list ) else [ servers ] <EOL> self . auth = ( username , password ) if username and password else None <EOL> self . timeout = timeout <EOL> def __repr__ ( self ) : <EOL> return '<STR_LIT>' % self . servers <EOL> @ staticmethod <EOL> def _parse_response ( response , clazz , is_list = False , resource_name = None ) : <EOL> """<STR_LIT>""" <EOL> target = response . json ( ) [ <EOL> resource_name ] if resource_name else response . json ( ) <EOL> if is_list : <EOL> return [ clazz . from_json ( resource ) for resource in target ] <EOL> else : <EOL> return clazz . from_json ( target ) <EOL> def _do_request ( self , method , path , params = None , data = None ) : <EOL> """<STR_LIT>""" <EOL> headers = { <EOL> '<STR_LIT:Content-Type>' : '<STR_LIT:application/json>' , '<STR_LIT>' : '<STR_LIT:application/json>' } <EOL> response = None <EOL> servers = list ( self . servers ) <EOL> while servers and response is None : <EOL> server = servers . pop ( <NUM_LIT:0> ) <EOL> url = '<STR_LIT>' . join ( [ server . rstrip ( '<STR_LIT:/>' ) , path ] ) <EOL> try : <EOL> response = requests . request ( <EOL> method , url , params = params , data = data , headers = headers , <EOL> auth = self . auth , timeout = self . timeout ) <EOL> marathon . log . info ( '<STR_LIT>' , server ) <EOL> except requests . exceptions . RequestException as e : <EOL> marathon . log . error ( <EOL> '<STR_LIT>' , url , e . message ) <EOL> if response is None : <EOL> raise MarathonError ( '<STR_LIT>' ) <EOL> if response . status_code >= <NUM_LIT> : <EOL> marathon . log . error ( '<STR_LIT>' . format ( <EOL> code = response . status_code , body = response . text ) ) <EOL> raise InternalServerError ( response ) <EOL> elif response . status_code >= <NUM_LIT> : <EOL> marathon . log . error ( '<STR_LIT>' . format ( <EOL> code = response . status_code , body = response . text ) ) <EOL> if response . status_code == <NUM_LIT> : <EOL> raise NotFoundError ( response ) <EOL> else : <EOL> raise MarathonHttpError ( response ) <EOL> elif response . status_code >= <NUM_LIT> : <EOL> marathon . log . warn ( '<STR_LIT>' . format ( <EOL> code = response . status_code , body = response . text ) ) <EOL> else : <EOL> marathon . log . debug ( '<STR_LIT>' . format ( <EOL> code = response . status_code , body = response . text ) ) <EOL> return response <EOL> def _do_sse_request ( self , path , params = None , data = None ) : <EOL> from sseclient import SSEClient <EOL> headers = { '<STR_LIT>' : '<STR_LIT>' } <EOL> messages = None <EOL> servers = list ( self . servers ) <EOL> while servers and messages is None : <EOL> server = servers . pop ( <NUM_LIT:0> ) <EOL> url = '<STR_LIT>' . join ( [ server . rstrip ( '<STR_LIT:/>' ) , path ] ) <EOL> try : <EOL> messages = SSEClient ( url , params = params , data = data , headers = headers , <EOL> auth = self . auth ) <EOL> except Exception as e : <EOL> marathon . log . error ( '<STR_LIT>' , url , e . message ) <EOL> if messages is None : <EOL> raise MarathonError ( '<STR_LIT>' ) <EOL> return messages <EOL> def list_endpoints ( self ) : <EOL> """<STR_LIT>""" <EOL> return MarathonEndpoint . from_tasks ( self . list_tasks ( ) ) <EOL> def create_app ( self , app_id , app ) : <EOL> """<STR_LIT>""" <EOL> app . id = app_id <EOL> data = app . to_json ( ) <EOL> response = self . _do_request ( '<STR_LIT:POST>' , '<STR_LIT>' , data = data ) <EOL> if response . status_code == <NUM_LIT> : <EOL> return self . _parse_response ( response , MarathonApp ) <EOL> else : <EOL> return False <EOL> def list_apps ( self , cmd = None , embed_tasks = False , <EOL> embed_failures = False , embed_task_stats = False , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> params = { } <EOL> if cmd : <EOL> params [ '<STR_LIT>' ] = cmd <EOL> if embed_failures : <EOL> params [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> elif embed_tasks : <EOL> params [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> elif embed_task_stats : <EOL> params [ '<STR_LIT>' ] = '<STR_LIT>' <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' , params = params ) <EOL> apps = self . _parse_response ( <EOL> response , MarathonApp , is_list = True , resource_name = '<STR_LIT>' ) <EOL> for k , v in kwargs . items ( ) : <EOL> apps = [ o for o in apps if getattr ( o , k ) == v ] <EOL> return apps <EOL> def get_app ( self , app_id , embed_tasks = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : '<STR_LIT>' } if embed_tasks else { } <EOL> response = self . _do_request ( <EOL> '<STR_LIT:GET>' , '<STR_LIT>' . format ( app_id = app_id ) , params = params ) <EOL> return self . _parse_response ( response , MarathonApp , resource_name = '<STR_LIT>' ) <EOL> def restart_app ( self , app_id , force = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : force } <EOL> response = self . _do_request ( <EOL> '<STR_LIT:POST>' , '<STR_LIT>' . format ( app_id = app_id ) , params = params ) <EOL> return response . json ( ) <EOL> def update_app ( self , app_id , app , force = False , minimal = True ) : <EOL> """<STR_LIT>""" <EOL> app . version = None <EOL> params = { '<STR_LIT>' : force } <EOL> data = app . to_json ( minimal = minimal ) <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( app_id = app_id ) , params = params , data = data ) <EOL> return response . json ( ) <EOL> def rollback_app ( self , app_id , version , force = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : force } <EOL> data = json . dumps ( { '<STR_LIT:version>' : version } ) <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( app_id = app_id ) , params = params , data = data ) <EOL> return response . json ( ) <EOL> def delete_app ( self , app_id , force = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : force } <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( app_id = app_id ) , params = params ) <EOL> return response . json ( ) <EOL> def scale_app ( self , app_id , instances = None , delta = None , force = False ) : <EOL> """<STR_LIT>""" <EOL> if instances is None and delta is None : <EOL> marathon . log . error ( '<STR_LIT>' ) <EOL> return <EOL> try : <EOL> app = self . get_app ( app_id ) <EOL> except NotFoundError : <EOL> marathon . log . error ( '<STR_LIT>' . format ( app = app_id ) ) <EOL> return <EOL> desired = instances if instances is not None else ( <EOL> app . instances + delta ) <EOL> return self . update_app ( app . id , MarathonApp ( instances = desired ) , force = force ) <EOL> def create_group ( self , group ) : <EOL> """<STR_LIT>""" <EOL> data = group . to_json ( ) <EOL> response = self . _do_request ( '<STR_LIT:POST>' , '<STR_LIT>' , data = data ) <EOL> return response . json ( ) <EOL> def list_groups ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> groups = self . _parse_response ( <EOL> response , MarathonGroup , is_list = True , resource_name = '<STR_LIT>' ) <EOL> for k , v in kwargs . items ( ) : <EOL> groups = [ o for o in groups if getattr ( o , k ) == v ] <EOL> return groups <EOL> def get_group ( self , group_id ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( <EOL> '<STR_LIT:GET>' , '<STR_LIT>' . format ( group_id = group_id ) ) <EOL> return self . _parse_response ( response , MarathonGroup , resource_name = '<STR_LIT>' ) <EOL> def update_group ( self , group_id , group , force = False , minimal = True ) : <EOL> """<STR_LIT>""" <EOL> group . version = None <EOL> params = { '<STR_LIT>' : force } <EOL> data = group . to_json ( minimal = minimal ) <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( group_id = group_id ) , data = data , params = params ) <EOL> return response . json ( ) <EOL> def rollback_group ( self , group_id , version , force = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : force } <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( group_id = group_id , <EOL> version = version ) , <EOL> params = params ) <EOL> return response . json ( ) <EOL> def delete_group ( self , group_id , force = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : force } <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( group_id = group_id ) , params = params ) <EOL> return response . json ( ) <EOL> def scale_group ( self , group_id , scale_by ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : scale_by } <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( group_id = group_id ) , params = params ) <EOL> return response . json ( ) <EOL> def list_tasks ( self , app_id = None , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> tasks = self . _parse_response ( <EOL> response , MarathonTask , is_list = True , resource_name = '<STR_LIT>' ) <EOL> if app_id : <EOL> tasks = [ <EOL> task for task in tasks if task . app_id . lstrip ( '<STR_LIT:/>' ) == app_id . lstrip ( '<STR_LIT:/>' ) ] <EOL> [ setattr ( t , '<STR_LIT>' , app_id ) <EOL> for t in tasks if app_id and t . app_id is None ] <EOL> for k , v in kwargs . items ( ) : <EOL> tasks = [ o for o in tasks if getattr ( o , k ) == v ] <EOL> return tasks <EOL> def kill_given_tasks ( self , task_ids , scale = False , force = None ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : scale } <EOL> if force is not None : <EOL> params [ '<STR_LIT>' ] = force <EOL> data = json . dumps ( { "<STR_LIT>" : task_ids } ) <EOL> response = self . _do_request ( <EOL> '<STR_LIT:POST>' , '<STR_LIT>' , params = params , data = data ) <EOL> return response == <NUM_LIT:200> <EOL> def kill_tasks ( self , app_id , scale = False , <EOL> host = None , batch_size = <NUM_LIT:0> , batch_delay = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> def batch ( iterable , size ) : <EOL> sourceiter = iter ( iterable ) <EOL> while True : <EOL> batchiter = itertools . islice ( sourceiter , size ) <EOL> yield itertools . chain ( [ next ( batchiter ) ] , batchiter ) <EOL> if batch_size == <NUM_LIT:0> : <EOL> params = { '<STR_LIT>' : scale } <EOL> if host : <EOL> params [ '<STR_LIT:host>' ] = host <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( app_id = app_id ) , params ) <EOL> if "<STR_LIT>" in response . json ( ) : <EOL> return self . _parse_response ( response , MarathonTask , is_list = True , resource_name = '<STR_LIT>' ) <EOL> else : <EOL> return response . json ( ) <EOL> else : <EOL> tasks = self . list_tasks ( <EOL> app_id , host = host ) if host else self . list_tasks ( app_id ) <EOL> for tbatch in batch ( tasks , batch_size ) : <EOL> killed_tasks = [ self . kill_task ( app_id , t . id , scale = scale ) <EOL> for t in tbatch ] <EOL> killed_task_ids = set ( t . id for t in killed_tasks ) <EOL> running_task_ids = killed_task_ids <EOL> while killed_task_ids . intersection ( running_task_ids ) : <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> running_task_ids = set ( <EOL> t . id for t in self . get_app ( app_id ) . tasks ) <EOL> if batch_delay == <NUM_LIT:0> : <EOL> desired_instances = self . get_app ( app_id ) . instances <EOL> running_instances = <NUM_LIT:0> <EOL> while running_instances < desired_instances : <EOL> time . sleep ( <NUM_LIT:1> ) <EOL> running_instances = sum ( <EOL> t . started_at is None for t in self . get_app ( app_id ) . tasks ) <EOL> else : <EOL> time . sleep ( batch_delay ) <EOL> return tasks <EOL> def kill_task ( self , app_id , task_id , scale = False ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : scale } <EOL> response = self . _do_request ( '<STR_LIT>' , '<STR_LIT>' <EOL> . format ( app_id = app_id , task_id = task_id ) , params ) <EOL> if "<STR_LIT>" in response . json ( ) : <EOL> return self . _parse_response ( response , MarathonTask , is_list = False , resource_name = '<STR_LIT>' ) <EOL> else : <EOL> return response . json ( ) <EOL> def list_versions ( self , app_id ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( <EOL> '<STR_LIT:GET>' , '<STR_LIT>' . format ( app_id = app_id ) ) <EOL> return [ version for version in response . json ( ) [ '<STR_LIT>' ] ] <EOL> def get_version ( self , app_id , version ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' <EOL> . format ( app_id = app_id , version = version ) ) <EOL> return MarathonApp ( response . json ( ) ) <EOL> def list_event_subscriptions ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return [ url for url in response . json ( ) [ '<STR_LIT>' ] ] <EOL> def create_event_subscription ( self , url ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : url } <EOL> response = self . _do_request ( '<STR_LIT:POST>' , '<STR_LIT>' , params ) <EOL> return response . json ( ) <EOL> def delete_event_subscription ( self , url ) : <EOL> """<STR_LIT>""" <EOL> params = { '<STR_LIT>' : url } <EOL> response = self . _do_request ( '<STR_LIT>' , '<STR_LIT>' , params ) <EOL> return response . json ( ) <EOL> def list_deployments ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return self . _parse_response ( response , MarathonDeployment , is_list = True ) <EOL> def list_queue ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return self . _parse_response ( response , MarathonQueueItem , is_list = True , resource_name = '<STR_LIT>' ) <EOL> def delete_deployment ( self , deployment_id , force = False ) : <EOL> """<STR_LIT>""" <EOL> if force : <EOL> params = { '<STR_LIT>' : True } <EOL> self . _do_request ( '<STR_LIT>' , '<STR_LIT>' . format ( <EOL> deployment = deployment_id ) , params = params ) <EOL> return { } <EOL> else : <EOL> response = self . _do_request ( <EOL> '<STR_LIT>' , '<STR_LIT>' . format ( deployment = deployment_id ) ) <EOL> return response . json ( ) <EOL> def get_info ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return self . _parse_response ( response , MarathonInfo ) <EOL> def get_leader ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return response . json ( ) <EOL> def delete_leader ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> return response . json ( ) <EOL> def ping ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return response . text <EOL> def get_metrics ( self ) : <EOL> """<STR_LIT>""" <EOL> response = self . _do_request ( '<STR_LIT:GET>' , '<STR_LIT>' ) <EOL> return response . json ( ) <EOL> def event_stream ( self ) : <EOL> """<STR_LIT>""" <EOL> messages = self . _do_sse_request ( '<STR_LIT>' ) <EOL> ef = EventFactory ( ) <EOL> for message in messages : <EOL> if not message . data : <EOL> continue <EOL> data = json . loads ( message . data ) <EOL> yield ef . process ( data ) </s>
