<s> from rexpro . exceptions import RexProConnectionException <EOL> from rexpro . messages import ErrorResponse <EOL> from contextlib import contextmanager <EOL> from hashlib import md5 <EOL> from Queue import Queue <EOL> import struct <EOL> from socket import socket <EOL> from textwrap import dedent <EOL> from rexpro import exceptions <EOL> from rexpro import messages <EOL> from rexpro import utils <EOL> class RexProSocket ( socket ) : <EOL> """<STR_LIT>""" <EOL> def send_message ( self , msg ) : <EOL> """<STR_LIT>""" <EOL> self . send ( msg . serialize ( ) ) <EOL> def get_response ( self ) : <EOL> """<STR_LIT>""" <EOL> msg_version = self . recv ( <NUM_LIT:1> ) <EOL> if not msg_version : <EOL> raise exceptions . RexProConnectionException ( '<STR_LIT>' ) <EOL> if bytearray ( [ msg_version ] ) [ <NUM_LIT:0> ] != <NUM_LIT:1> : <EOL> raise exceptions . RexProConnectionException ( '<STR_LIT>' . format ( ) ) <EOL> serializer_type = self . recv ( <NUM_LIT:1> ) <EOL> if bytearray ( serializer_type ) [ <NUM_LIT:0> ] != <NUM_LIT:0> : <EOL> raise exceptions . RexProConnectionException ( '<STR_LIT>' . format ( ) ) <EOL> self . recv ( <NUM_LIT:4> ) <EOL> msg_type = self . recv ( <NUM_LIT:1> ) <EOL> msg_type = bytearray ( msg_type ) [ <NUM_LIT:0> ] <EOL> msg_len = struct . unpack ( '<STR_LIT>' , self . recv ( <NUM_LIT:4> ) ) [ <NUM_LIT:0> ] <EOL> response = '<STR_LIT>' <EOL> while len ( response ) < msg_len : <EOL> response += self . recv ( msg_len ) <EOL> MessageTypes = messages . MessageTypes <EOL> type_map = { <EOL> MessageTypes . ERROR : messages . ErrorResponse , <EOL> MessageTypes . SESSION_RESPONSE : messages . SessionResponse , <EOL> MessageTypes . SCRIPT_RESPONSE : messages . MsgPackScriptResponse <EOL> } <EOL> if msg_type not in type_map : <EOL> raise exceptions . RexProConnectionException ( "<STR_LIT>" . format ( msg_type ) ) <EOL> return type_map [ msg_type ] . deserialize ( response ) <EOL> class RexProConnectionPool ( object ) : <EOL> def __init__ ( self , host , port , size ) : <EOL> """<STR_LIT>""" <EOL> self . host = host <EOL> self . port = port <EOL> self . size = size <EOL> self . pool = Queue ( ) <EOL> for i in range ( size ) : <EOL> self . pool . put ( self . _new_conn ( ) ) <EOL> def _new_conn ( self ) : <EOL> """<STR_LIT>""" <EOL> conn = RexProSocket ( ) <EOL> conn . connect ( ( self . host , self . port ) ) <EOL> return conn <EOL> def get ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . pool . empty ( ) : <EOL> return self . _new_conn ( ) <EOL> return self . pool . get ( ) <EOL> def put ( self , conn ) : <EOL> """<STR_LIT>""" <EOL> if self . pool . qsize ( ) >= self . size : <EOL> conn . close ( ) <EOL> return <EOL> self . pool . put ( conn ) <EOL> @ contextmanager <EOL> def contextual_connection ( self ) : <EOL> """<STR_LIT>""" <EOL> conn = self . get ( ) <EOL> yield conn <EOL> self . put ( conn ) <EOL> def __del__ ( self ) : <EOL> while not self . pool . empty ( ) : <EOL> self . pool . get ( ) . close ( ) <EOL> class RexProConnection ( object ) : <EOL> def __init__ ( self , host , port , graph_name , graph_obj_name = '<STR_LIT:g>' , username = '<STR_LIT>' , password = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> self . host = host <EOL> self . port = port <EOL> self . graph_name = graph_name <EOL> self . username = username <EOL> self . password = password <EOL> self . graph_features = None <EOL> self . _conn = RexProSocket ( ) <EOL> self . _conn . connect ( ( self . host , self . port ) ) <EOL> self . _in_transaction = False <EOL> self . _session_key = None <EOL> self . _open_session ( ) <EOL> def _open_session ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _conn . send_message ( <EOL> messages . SessionRequest ( <EOL> username = self . username , <EOL> password = self . password , <EOL> graph_name = self . graph_name <EOL> ) <EOL> ) <EOL> response = self . _conn . get_response ( ) <EOL> if isinstance ( response , ErrorResponse ) : <EOL> raise RexProConnectionException ( response . message ) <EOL> self . _session_key = response . session_key <EOL> self . graph_features = self . execute ( '<STR_LIT>' ) <EOL> def open_transaction ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _in_transaction : <EOL> raise exceptions . RexProScriptException ( "<STR_LIT>" ) <EOL> self . execute ( <EOL> script = '<STR_LIT>' , <EOL> isolate = False <EOL> ) <EOL> self . _in_transaction = True <EOL> def close_transaction ( self , success = True ) : <EOL> """<STR_LIT>""" <EOL> if not self . _in_transaction : <EOL> raise exceptions . RexProScriptException ( "<STR_LIT>" ) <EOL> self . execute ( <EOL> script = '<STR_LIT>' . format ( '<STR_LIT>' if success else '<STR_LIT>' ) , <EOL> isolate = False <EOL> ) <EOL> self . _in_transaction = False <EOL> def close ( self ) : <EOL> self . _conn . send_message ( <EOL> messages . SessionRequest ( <EOL> session_key = self . _session_key , <EOL> graph_name = self . graph_name , <EOL> kill_session = True <EOL> ) <EOL> ) <EOL> response = self . _conn . get_response ( ) <EOL> if isinstance ( response , ErrorResponse ) : <EOL> raise RexProConnectionException ( response . message ) <EOL> @ contextmanager <EOL> def transaction ( self ) : <EOL> """<STR_LIT>""" <EOL> self . open_transaction ( ) <EOL> yield <EOL> self . close_transaction ( ) <EOL> def execute ( self , script , params = { } , isolate = True , transaction = True , pretty = False ) : <EOL> """<STR_LIT>""" <EOL> self . _conn . send_message ( <EOL> messages . ScriptRequest ( <EOL> script = script , <EOL> params = params , <EOL> session_key = self . _session_key , <EOL> isolate = isolate , <EOL> in_transaction = transaction <EOL> ) <EOL> ) <EOL> response = self . _conn . get_response ( ) <EOL> if isinstance ( response , messages . ErrorResponse ) : <EOL> raise exceptions . RexProScriptException ( response . message ) <EOL> return response . results </s>
