<s> import Queue <EOL> import logging <EOL> import socket <EOL> import threading <EOL> import time <EOL> import sasl <EOL> import sys <EOL> from thrift . Thrift import TType , TApplicationException <EOL> from thrift . transport . TSocket import TSocket <EOL> from thrift . transport . TSSLSocket import TSSLSocket <EOL> from thrift . transport . TTransport import TBufferedTransport , TFramedTransport , TMemoryBuffer , TTransportException <EOL> from thrift . protocol . TBinaryProtocol import TBinaryProtocol <EOL> from thrift . protocol . TMultiplexedProtocol import TMultiplexedProtocol <EOL> from desktop . lib . python_util import create_synchronous_io_multiplexer <EOL> from desktop . lib . thrift_ . http_client import THttpClient <EOL> from desktop . lib . thrift_sasl import TSaslClientTransport <EOL> from desktop . lib . exceptions import StructuredException , StructuredThriftTransportException <EOL> MAX_RECURSION_DEPTH = <NUM_LIT:50> <EOL> WARN_LEVEL_CALL_DURATION_MS = <NUM_LIT> <EOL> INFO_LEVEL_CALL_DURATION_MS = <NUM_LIT:1000> <EOL> class LifoQueue ( Queue . Queue ) : <EOL> '''<STR_LIT>''' <EOL> def _init ( self , maxsize ) : <EOL> self . queue = [ ] <EOL> self . maxsize = maxsize <EOL> def _qsize ( self , len = len ) : <EOL> return len ( self . queue ) <EOL> def _put ( self , item ) : <EOL> self . queue . append ( item ) <EOL> def _get ( self ) : <EOL> return self . queue . pop ( ) <EOL> class ConnectionConfig ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , klass , host , port , service_name , <EOL> use_sasl = False , <EOL> use_ssl = False , <EOL> kerberos_principal = "<STR_LIT>" , <EOL> mechanism = '<STR_LIT>' , <EOL> username = '<STR_LIT>' , <EOL> password = '<STR_LIT>' , <EOL> ca_certs = None , <EOL> keyfile = None , <EOL> certfile = None , <EOL> validate = False , <EOL> timeout_seconds = <NUM_LIT> , <EOL> transport = '<STR_LIT>' , <EOL> multiple = False , <EOL> transport_mode = '<STR_LIT>' , <EOL> http_url = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> self . klass = klass <EOL> self . host = host <EOL> self . port = port <EOL> self . service_name = service_name <EOL> self . use_sasl = use_sasl <EOL> self . use_ssl = use_ssl <EOL> self . mechanism = mechanism <EOL> self . username = username <EOL> self . password = password <EOL> self . kerberos_principal = kerberos_principal <EOL> self . ca_certs = ca_certs <EOL> self . keyfile = keyfile <EOL> self . certfile = certfile <EOL> self . validate = validate <EOL> self . timeout_seconds = timeout_seconds <EOL> self . transport = transport <EOL> self . multiple = multiple <EOL> self . transport_mode = transport_mode <EOL> self . http_url = http_url <EOL> def __str__ ( self ) : <EOL> return '<STR_LIT:U+002CU+0020>' . join ( map ( str , [ self . klass , self . host , self . port , self . service_name , self . use_sasl , self . kerberos_principal , self . timeout_seconds , <EOL> self . mechanism , self . username , self . use_ssl , self . ca_certs , self . keyfile , self . certfile , self . validate , self . transport , <EOL> self . multiple , self . transport_mode , self . http_url ] ) ) <EOL> class ConnectionPooler ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , poolsize = <NUM_LIT:10> ) : <EOL> self . pooldict = { } <EOL> self . poolsize = poolsize <EOL> self . dictlock = threading . Lock ( ) <EOL> def get_client ( self , conf , get_client_timeout = None ) : <EOL> """<STR_LIT>""" <EOL> if _get_pool_key ( conf ) not in self . pooldict : <EOL> self . dictlock . acquire ( ) <EOL> try : <EOL> if _get_pool_key ( conf ) not in self . pooldict : <EOL> q = LifoQueue ( self . poolsize ) <EOL> for i in xrange ( self . poolsize ) : <EOL> client = construct_superclient ( conf ) <EOL> client . CID = i <EOL> q . put ( client , False ) <EOL> self . pooldict [ _get_pool_key ( conf ) ] = q <EOL> finally : <EOL> self . dictlock . release ( ) <EOL> connection = None <EOL> start_pool_get_time = time . time ( ) <EOL> has_waited_for = <NUM_LIT:0> <EOL> while connection is None : <EOL> if get_client_timeout is not None : <EOL> this_round_timeout = max ( min ( get_client_timeout - has_waited_for , <NUM_LIT:1> ) , <NUM_LIT:0> ) <EOL> else : <EOL> this_round_timeout = None <EOL> try : <EOL> connection = self . pooldict [ _get_pool_key ( conf ) ] . get ( <EOL> block = True , timeout = this_round_timeout ) <EOL> except Queue . Empty : <EOL> has_waited_for = time . time ( ) - start_pool_get_time <EOL> if get_client_timeout is not None and has_waited_for > get_client_timeout : <EOL> raise socket . timeout ( <EOL> ( "<STR_LIT>" + <EOL> "<STR_LIT>" ) % ( has_waited_for , conf . service_name ) ) <EOL> logging . warn ( "<STR_LIT>" % <EOL> ( has_waited_for , conf . host , conf . port ) ) <EOL> return connection <EOL> def return_client ( self , conf , client ) : <EOL> """<STR_LIT>""" <EOL> self . pooldict [ _get_pool_key ( conf ) ] . put ( client ) <EOL> def _get_pool_key ( conf ) : <EOL> """<STR_LIT>""" <EOL> return ( conf . klass , conf . host , conf . port ) <EOL> def construct_superclient ( conf ) : <EOL> """<STR_LIT>""" <EOL> service , protocol , transport = connect_to_thrift ( conf ) <EOL> return SuperClient ( service , transport , timeout_seconds = conf . timeout_seconds ) <EOL> def connect_to_thrift ( conf ) : <EOL> """<STR_LIT>""" <EOL> if conf . transport_mode == '<STR_LIT:http>' : <EOL> mode = THttpClient ( conf . http_url ) <EOL> else : <EOL> if conf . use_ssl : <EOL> mode = TSSLSocket ( conf . host , conf . port , validate = conf . validate , ca_certs = conf . ca_certs , keyfile = conf . keyfile , certfile = conf . certfile ) <EOL> else : <EOL> mode = TSocket ( conf . host , conf . port ) <EOL> if conf . timeout_seconds : <EOL> mode . setTimeout ( conf . timeout_seconds * <NUM_LIT> ) <EOL> if conf . transport_mode == '<STR_LIT:http>' : <EOL> if conf . use_sasl and conf . mechanism != '<STR_LIT>' : <EOL> mode . set_kerberos_auth ( ) <EOL> else : <EOL> mode . set_basic_auth ( conf . username , conf . password ) <EOL> if conf . transport_mode == '<STR_LIT>' and conf . use_sasl : <EOL> def sasl_factory ( ) : <EOL> saslc = sasl . Client ( ) <EOL> saslc . setAttr ( "<STR_LIT:host>" , str ( conf . host ) ) <EOL> saslc . setAttr ( "<STR_LIT>" , str ( conf . kerberos_principal ) ) <EOL> if conf . mechanism == '<STR_LIT>' : <EOL> saslc . setAttr ( "<STR_LIT:username>" , str ( conf . username ) ) <EOL> saslc . setAttr ( "<STR_LIT:password>" , str ( conf . password ) ) <EOL> saslc . init ( ) <EOL> return saslc <EOL> transport = TSaslClientTransport ( sasl_factory , conf . mechanism , mode ) <EOL> elif conf . transport == '<STR_LIT>' : <EOL> transport = TFramedTransport ( mode ) <EOL> else : <EOL> transport = TBufferedTransport ( mode ) <EOL> protocol = TBinaryProtocol ( transport ) <EOL> if conf . multiple : <EOL> protocol = TMultiplexedProtocol ( protocol , conf . service_name ) <EOL> service = conf . klass ( protocol ) <EOL> return service , protocol , transport <EOL> _connection_pool = ConnectionPooler ( ) <EOL> def get_client ( klass , host , port , service_name , ** kwargs ) : <EOL> conf = ConnectionConfig ( klass , host , port , service_name , ** kwargs ) <EOL> return PooledClient ( conf ) <EOL> def _grab_transport_from_wrapper ( outer_transport ) : <EOL> if isinstance ( outer_transport , TBufferedTransport ) : <EOL> return outer_transport . _TBufferedTransport__trans <EOL> elif isinstance ( outer_transport , TSaslClientTransport ) : <EOL> return outer_transport . _trans <EOL> elif isinstance ( outer_transport , TFramedTransport ) : <EOL> return outer_transport . _TFramedTransport__trans <EOL> else : <EOL> raise Exception ( "<STR_LIT>" + outer_transport . __class__ ) <EOL> class PooledClient ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , conf ) : <EOL> self . conf = conf <EOL> def __getattr__ ( self , attr_name ) : <EOL> if attr_name in self . __dict__ : <EOL> return self . __dict__ [ attr_name ] <EOL> superclient = _connection_pool . get_client ( self . conf , <EOL> get_client_timeout = self . conf . timeout_seconds ) <EOL> try : <EOL> attr = getattr ( superclient , attr_name ) <EOL> if callable ( attr ) : <EOL> return self . _wrap_callable ( attr_name ) <EOL> else : <EOL> return attr <EOL> finally : <EOL> self . _return_client ( superclient ) <EOL> def _wrap_callable ( self , attr_name ) : <EOL> def wrapper ( * args , ** kwargs ) : <EOL> superclient = _connection_pool . get_client ( self . conf ) <EOL> try : <EOL> attr = getattr ( superclient , attr_name ) <EOL> try : <EOL> sock = self . conf . transport_mode != '<STR_LIT:http>' and _grab_transport_from_wrapper ( superclient . transport ) . handle <EOL> if sock and create_synchronous_io_multiplexer ( ) . read ( [ sock ] ) : <EOL> superclient . transport . close ( ) <EOL> superclient . transport . open ( ) <EOL> superclient . set_timeout ( self . conf . timeout_seconds ) <EOL> return attr ( * args , ** kwargs ) <EOL> except TApplicationException , e : <EOL> logging . info ( "<STR_LIT>" + str ( e ) , exc_info = False ) <EOL> raise StructuredException ( '<STR_LIT>' , str ( e ) , data = None , error_code = <NUM_LIT> ) <EOL> except socket . error , e : <EOL> logging . info ( "<STR_LIT>" + str ( e ) , exc_info = False ) <EOL> raise StructuredException ( '<STR_LIT>' , str ( e ) , data = None , error_code = <NUM_LIT> ) <EOL> except TTransportException , e : <EOL> logging . info ( "<STR_LIT>" + str ( e ) , exc_info = False ) <EOL> raise StructuredThriftTransportException ( e , error_code = <NUM_LIT> ) <EOL> except Exception , e : <EOL> logging . info ( "<STR_LIT>" + str ( e ) , exc_info = False ) <EOL> raise <EOL> finally : <EOL> self . _return_client ( superclient ) <EOL> wrapper . attr = attr_name <EOL> return wrapper <EOL> def _return_client ( self , superclient ) : <EOL> _connection_pool . return_client ( self . conf , superclient ) <EOL> class SuperClient ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , wrapped_client , transport , timeout_seconds = None ) : <EOL> self . wrapped = wrapped_client <EOL> self . transport = transport <EOL> self . timeout_seconds = timeout_seconds <EOL> def __getattr__ ( self , attr ) : <EOL> if attr in self . __dict__ : <EOL> return self . __dict__ [ attr ] <EOL> res = getattr ( self . wrapped , attr ) <EOL> if not hasattr ( res , '<STR_LIT>' ) : <EOL> return res <EOL> def wrapper ( * args , ** kwargs ) : <EOL> tries_left = <NUM_LIT:3> <EOL> while tries_left : <EOL> sys . exc_clear ( ) <EOL> try : <EOL> if not self . transport . isOpen ( ) : <EOL> self . transport . open ( ) <EOL> st = time . time ( ) <EOL> logging . debug ( "<STR_LIT>" <EOL> % ( str ( self . wrapped . __class__ ) , attr , repr ( args ) , repr ( kwargs ) ) ) <EOL> ret = res ( * args , ** kwargs ) <EOL> log_msg = repr ( ret ) <EOL> if len ( log_msg ) > <NUM_LIT:1000> : <EOL> log_msg = log_msg [ <NUM_LIT:0> : <NUM_LIT:1000> ] + "<STR_LIT>" <EOL> duration = time . time ( ) - st <EOL> logmsg = "<STR_LIT>" % ( <EOL> str ( self . wrapped . __class__ ) , attr , duration * <NUM_LIT:1000> , log_msg ) <EOL> if duration >= WARN_LEVEL_CALL_DURATION_MS : <EOL> logging . warn ( logmsg ) <EOL> elif duration >= INFO_LEVEL_CALL_DURATION_MS : <EOL> logging . info ( logmsg ) <EOL> else : <EOL> logging . debug ( logmsg ) <EOL> return ret <EOL> except socket . error , e : <EOL> pass <EOL> except TTransportException , e : <EOL> pass <EOL> except Exception , e : <EOL> logging . exception ( "<STR_LIT>" ) <EOL> raise <EOL> self . transport . close ( ) <EOL> if isinstance ( e , socket . timeout ) : <EOL> logging . warn ( "<STR_LIT>" % attr ) <EOL> raise <EOL> else : <EOL> tries_left -= <NUM_LIT:1> <EOL> if tries_left : <EOL> logging . info ( "<STR_LIT>" + str ( e ) , exc_info = <NUM_LIT:0> ) <EOL> logging . warn ( "<STR_LIT>" + attr ) <EOL> raise <EOL> return wrapper <EOL> def set_timeout ( self , timeout_seconds ) : <EOL> if timeout_seconds != self . timeout_seconds : <EOL> self . timeout_seconds = timeout_seconds <EOL> if self . timeout_seconds is not None : <EOL> _grab_transport_from_wrapper ( self . transport ) . setTimeout ( self . timeout_seconds * <NUM_LIT:1000> ) <EOL> else : <EOL> _grab_transport_from_wrapper ( self . transport ) . setTimeout ( None ) <EOL> def simpler_string ( thrift_obj ) : <EOL> """<STR_LIT>""" <EOL> L = [ ] <EOL> for key , value in thrift_obj . __dict__ . iteritems ( ) : <EOL> if value is None : <EOL> continue <EOL> if hasattr ( value , "<STR_LIT>" ) : <EOL> L . append ( "<STR_LIT>" % ( key , simpler_string ( value ) ) ) <EOL> else : <EOL> L . append ( "<STR_LIT>" % ( key , value ) ) <EOL> return '<STR_LIT>' % ( thrift_obj . __class__ . __name__ , '<STR_LIT:U+002CU+0020>' . join ( L ) ) <EOL> def from_bytes ( klass , data ) : <EOL> """<STR_LIT>""" <EOL> obj = klass ( ) <EOL> b = TMemoryBuffer ( data ) <EOL> p = TBinaryProtocol ( b ) <EOL> obj . read ( p ) <EOL> return obj <EOL> def to_bytes ( obj ) : <EOL> """<STR_LIT>""" <EOL> b = TMemoryBuffer ( ) <EOL> p = TBinaryProtocol ( b ) <EOL> obj . write ( p ) <EOL> return b . getvalue ( ) <EOL> def thrift2json ( tft ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( tft , type ( None ) ) : <EOL> return None <EOL> if isinstance ( tft , ( float , int , long , complex , basestring ) ) : <EOL> return tft <EOL> if isinstance ( tft , dict ) : <EOL> d = { } <EOL> for key , val in tft . iteritems ( ) : <EOL> d [ key ] = thrift2json ( val ) <EOL> return d <EOL> if isinstance ( tft , list ) : <EOL> return [ thrift2json ( x ) for x in tft ] <EOL> if isinstance ( tft , set ) : <EOL> return dict ( ( x , True ) for x in tft ) <EOL> json = { } <EOL> d = { } <EOL> if hasattr ( tft , "<STR_LIT>" ) : <EOL> d = tft . __dict__ <EOL> else : <EOL> if hasattr ( tft , "<STR_LIT>" ) : <EOL> d = tft . __slots__ <EOL> else : <EOL> return { } <EOL> for k in d : <EOL> v = getattr ( tft , k ) <EOL> json [ k ] = thrift2json ( v ) <EOL> return json <EOL> def _jsonable2thrift_helper ( jsonable , type_enum , spec_args , default , recursion_depth = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> if recursion_depth > MAX_RECURSION_DEPTH : <EOL> raise Exception ( "<STR_LIT>" ) <EOL> if jsonable is None : <EOL> return default <EOL> def check_bits ( jsonable , bits ) : <EOL> """<STR_LIT>""" <EOL> check_type ( jsonable , ( int , long ) ) <EOL> min_val = - <NUM_LIT:1> << ( bits - <NUM_LIT:1> ) <EOL> max_val = ( <NUM_LIT:1> << ( bits - <NUM_LIT:1> ) ) - <NUM_LIT:1> <EOL> assert min_val <= jsonable , "<STR_LIT>" % ( jsonable , min_val ) <EOL> assert max_val >= jsonable , "<STR_LIT>" % ( jsonable , max_val ) <EOL> def check_type ( jsonable , expected ) : <EOL> assert isinstance ( jsonable , expected ) , "<STR_LIT>" % ( repr ( jsonable ) , repr ( expected ) ) <EOL> if type_enum == TType . BOOL : <EOL> check_type ( jsonable , bool ) <EOL> return jsonable <EOL> elif type_enum == TType . BYTE : <EOL> check_bits ( jsonable , <NUM_LIT:8> ) <EOL> return jsonable <EOL> elif type_enum == TType . I16 : <EOL> check_bits ( jsonable , <NUM_LIT:16> ) <EOL> return jsonable <EOL> elif type_enum == TType . I32 : <EOL> check_bits ( jsonable , <NUM_LIT:32> ) <EOL> return jsonable <EOL> elif type_enum == TType . I64 : <EOL> check_bits ( jsonable , <NUM_LIT:64> ) <EOL> return jsonable <EOL> elif type_enum == TType . DOUBLE : <EOL> check_type ( jsonable , float ) <EOL> return jsonable <EOL> elif type_enum == TType . STRING : <EOL> assert isinstance ( jsonable , basestring ) <EOL> return jsonable <EOL> elif type_enum == TType . STRUCT : <EOL> check_type ( jsonable , dict ) <EOL> thrift_type , thrift_spec = spec_args <EOL> out = thrift_type ( ) <EOL> for spec in thrift_spec : <EOL> if spec is None : <EOL> continue <EOL> _ , cur_type_enum , cur_name , cur_spec_args , cur_default = spec <EOL> value = _jsonable2thrift_helper ( jsonable . get ( cur_name ) , <EOL> cur_type_enum , cur_spec_args , cur_default , recursion_depth + <NUM_LIT:1> ) <EOL> setattr ( out , cur_name , value ) <EOL> return out <EOL> elif type_enum == TType . MAP : <EOL> check_type ( jsonable , dict ) <EOL> key_type_enum , key_spec_args , val_type_enum , val_spec_args = spec_args <EOL> out = dict ( ) <EOL> for k_jsonable , v_jsonable in jsonable . iteritems ( ) : <EOL> k = _jsonable2thrift_helper ( k_jsonable , key_type_enum , key_spec_args , None , recursion_depth + <NUM_LIT:1> ) <EOL> v = _jsonable2thrift_helper ( v_jsonable , val_type_enum , val_spec_args , None , recursion_depth + <NUM_LIT:1> ) <EOL> out [ k ] = v <EOL> return out <EOL> elif type_enum == TType . SET : <EOL> set_type_enum , set_spec_args = spec_args <EOL> out = set ( ) <EOL> for k , v in jsonable . iteritems ( ) : <EOL> assert v is True , "<STR_LIT>" % repr ( v ) <EOL> out . add ( _jsonable2thrift_helper ( k , set_type_enum , set_spec_args , None , recursion_depth + <NUM_LIT:1> ) ) <EOL> return out <EOL> elif type_enum == TType . LIST : <EOL> check_type ( jsonable , list ) <EOL> list_type_enum , list_spec_args = spec_args <EOL> out = list ( ) <EOL> for x in jsonable : <EOL> out . append ( _jsonable2thrift_helper ( x , list_type_enum , list_spec_args , None , recursion_depth + <NUM_LIT:1> ) ) <EOL> return out <EOL> else : <EOL> raise Exception ( "<STR_LIT>" % ( repr ( type_enum ) , repr ( jsonable ) ) ) <EOL> def jsonable2thrift ( jsonable , thrift_class ) : <EOL> """<STR_LIT>""" <EOL> return _jsonable2thrift_helper ( <EOL> jsonable , <EOL> TType . STRUCT , <EOL> ( thrift_class , thrift_class . thrift_spec ) , <EOL> default = None , <EOL> recursion_depth = <NUM_LIT:0> <EOL> ) <EOL> def enum_as_sequence ( enum ) : <EOL> """<STR_LIT>""" <EOL> return filter ( lambda x : not x . startswith ( "<STR_LIT>" ) <EOL> and x not in [ "<STR_LIT>" , "<STR_LIT>" ] , dir ( enum ) ) <EOL> def fixup_enums ( obj , name_class_map , suffix = "<STR_LIT>" ) : <EOL> """<STR_LIT>""" <EOL> for n in name_class_map . keys ( ) : <EOL> c = name_class_map [ n ] <EOL> setattr ( obj , n + suffix , c . _VALUES_TO_NAMES [ getattr ( obj , n ) ] ) <EOL> return obj <EOL> def is_thrift_struct ( o ) : <EOL> return hasattr ( o . __class__ , "<STR_LIT>" ) </s>
