<s> """<STR_LIT>""" <EOL> __author__ = '<STR_LIT>' <EOL> import contextlib <EOL> import gevent <EOL> from gevent . queue import Queue <EOL> from gevent . socket import wait_read , wait_write <EOL> import time <EOL> import sys <EOL> import simplejson as json <EOL> try : <EOL> import psycopg2 <EOL> from psycopg2 import OperationalError , ProgrammingError , DatabaseError , IntegrityError , extensions <EOL> from psycopg2 . extensions import ISOLATION_LEVEL_AUTOCOMMIT <EOL> from psycopg2 . extensions import connection as _connection <EOL> from psycopg2 . extensions import cursor as _cursor <EOL> from psycopg2 . extras import register_default_json <EOL> except ImportError : <EOL> print "<STR_LIT>" <EOL> def gevent_wait_callback ( conn , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> while <NUM_LIT:1> : <EOL> state = conn . poll ( ) <EOL> if state == extensions . POLL_OK : <EOL> break <EOL> elif state == extensions . POLL_READ : <EOL> wait_read ( conn . fileno ( ) , timeout = timeout ) <EOL> elif state == extensions . POLL_WRITE : <EOL> wait_write ( conn . fileno ( ) , timeout = timeout ) <EOL> else : <EOL> raise OperationalError ( <EOL> "<STR_LIT>" % state ) <EOL> extensions . set_wait_callback ( gevent_wait_callback ) <EOL> register_default_json ( None , globally = True , loads = json . loads ) <EOL> class DatabaseConnectionPool ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , maxsize = <NUM_LIT:100> ) : <EOL> if not isinstance ( maxsize , ( int , long ) ) : <EOL> raise TypeError ( '<STR_LIT>' % ( maxsize , ) ) <EOL> self . maxsize = maxsize <EOL> self . pool = Queue ( ) <EOL> self . size = <NUM_LIT:0> <EOL> def get ( self ) : <EOL> pool = self . pool <EOL> if self . size >= self . maxsize or pool . qsize ( ) : <EOL> return pool . get ( ) <EOL> else : <EOL> self . size += <NUM_LIT:1> <EOL> try : <EOL> new_item = self . create_connection ( ) <EOL> except : <EOL> self . size -= <NUM_LIT:1> <EOL> raise <EOL> return new_item <EOL> def put ( self , item ) : <EOL> self . pool . put ( item ) <EOL> def closeall ( self ) : <EOL> while not self . pool . empty ( ) : <EOL> conn = self . pool . get_nowait ( ) <EOL> try : <EOL> conn . close ( ) <EOL> self . size -= <NUM_LIT:1> <EOL> except Exception : <EOL> pass <EOL> @ contextlib . contextmanager <EOL> def connection ( self , isolation_level = None ) : <EOL> conn = self . get ( ) <EOL> try : <EOL> if isolation_level is not None : <EOL> if conn . isolation_level == isolation_level : <EOL> isolation_level = None <EOL> else : <EOL> conn . set_isolation_level ( isolation_level ) <EOL> yield conn <EOL> except : <EOL> if conn . closed : <EOL> conn = None <EOL> self . closeall ( ) <EOL> else : <EOL> conn = self . _rollback ( conn ) <EOL> raise <EOL> else : <EOL> if conn . closed : <EOL> raise OperationalError ( "<STR_LIT>" % ( conn , ) ) <EOL> conn . commit ( ) <EOL> finally : <EOL> if conn is not None and not conn . closed : <EOL> if isolation_level is not None : <EOL> conn . set_isolation_level ( isolation_level ) <EOL> self . put ( conn ) <EOL> @ contextlib . contextmanager <EOL> def cursor ( self , * args , ** kwargs ) : <EOL> isolation_level = kwargs . pop ( '<STR_LIT>' , None ) <EOL> conn = self . get ( ) <EOL> try : <EOL> if isolation_level is not None : <EOL> if conn . isolation_level == isolation_level : <EOL> isolation_level = None <EOL> else : <EOL> conn . set_isolation_level ( isolation_level ) <EOL> tracer = kwargs . pop ( "<STR_LIT>" , None ) <EOL> cur = conn . cursor ( * args , ** kwargs ) <EOL> if isinstance ( cur , TracingCursor ) : <EOL> cur . _tracer = tracer <EOL> yield cur <EOL> except : <EOL> if conn . closed : <EOL> conn = None <EOL> self . closeall ( ) <EOL> else : <EOL> conn = self . _rollback ( conn ) <EOL> raise <EOL> else : <EOL> if conn . closed : <EOL> raise OperationalError ( "<STR_LIT>" % ( conn , ) ) <EOL> conn . commit ( ) <EOL> finally : <EOL> if conn is not None and not conn . closed : <EOL> if isolation_level is not None : <EOL> conn . set_isolation_level ( isolation_level ) <EOL> self . put ( conn ) <EOL> def _rollback ( self , conn ) : <EOL> try : <EOL> conn . rollback ( ) <EOL> except : <EOL> gevent . get_hub ( ) . handle_error ( conn , * sys . exc_info ( ) ) <EOL> return <EOL> return conn <EOL> def execute ( self , * args , ** kwargs ) : <EOL> with self . cursor ( ** kwargs ) as cursor : <EOL> cursor . execute ( * args ) <EOL> return cursor . rowcount <EOL> def fetchone ( self , * args , ** kwargs ) : <EOL> with self . cursor ( ** kwargs ) as cursor : <EOL> cursor . execute ( * args ) <EOL> return cursor . fetchone ( ) <EOL> def fetchall ( self , * args , ** kwargs ) : <EOL> with self . cursor ( ** kwargs ) as cursor : <EOL> cursor . execute ( * args ) <EOL> return cursor . fetchall ( ) <EOL> def fetchiter ( self , * args , ** kwargs ) : <EOL> with self . cursor ( ** kwargs ) as cursor : <EOL> cursor . execute ( * args ) <EOL> while True : <EOL> items = cursor . fetchmany ( ) <EOL> if not items : <EOL> break <EOL> for item in items : <EOL> yield item <EOL> class PostgresConnectionPool ( DatabaseConnectionPool ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . connect = kwargs . pop ( '<STR_LIT>' , psycopg2 . connect ) <EOL> self . tracer = kwargs . pop ( '<STR_LIT>' , None ) <EOL> maxsize = kwargs . pop ( '<STR_LIT>' , None ) <EOL> self . args = args <EOL> self . kwargs = kwargs <EOL> if self . tracer : <EOL> self . kwargs . setdefault ( "<STR_LIT>" , TracingConnection ) <EOL> DatabaseConnectionPool . __init__ ( self , maxsize ) <EOL> def create_connection ( self ) : <EOL> conn = self . connect ( * self . args , ** self . kwargs ) <EOL> if self . tracer : <EOL> conn . set_tracer ( self . tracer ) <EOL> return conn <EOL> def psycopg2_connect ( dsn = None , * args , ** kwargs ) : <EOL> if dsn is None : <EOL> c_host = kwargs . pop ( "<STR_LIT>" , None ) or "<STR_LIT:localhost>" <EOL> c_port = kwargs . pop ( "<STR_LIT>" , None ) or "<STR_LIT>" <EOL> c_dbname = kwargs . pop ( "<STR_LIT>" , None ) or "<STR_LIT>" <EOL> c_user = kwargs . pop ( "<STR_LIT>" , None ) or "<STR_LIT>" <EOL> c_password = kwargs . pop ( "<STR_LIT>" , None ) or "<STR_LIT>" <EOL> dsn = "<STR_LIT>" % ( c_host , c_port , c_dbname , c_user , c_password ) <EOL> tracer = kwargs . pop ( "<STR_LIT>" , None ) <EOL> trace_stmt = kwargs . pop ( "<STR_LIT>" , None ) <EOL> if tracer : <EOL> conn = psycopg2 . connect ( dsn , * args , connection_factory = TracingConnection , ** kwargs ) <EOL> conn . set_tracer ( tracer , trace_stmt ) <EOL> else : <EOL> conn = psycopg2 . connect ( dsn , * args , ** kwargs ) <EOL> return conn <EOL> class TracingConnection ( _connection ) : <EOL> """<STR_LIT>""" <EOL> def set_tracer ( self , tracer , trace_stmt = None ) : <EOL> self . _tracer = tracer <EOL> self . _trace_stmt = trace_stmt <EOL> def cursor ( self , * args , ** kwargs ) : <EOL> kwargs . setdefault ( '<STR_LIT>' , TracingCursor ) <EOL> return super ( TracingConnection , self ) . cursor ( * args , ** kwargs ) <EOL> class TracingCursor ( _cursor ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . _tracer = kwargs . pop ( "<STR_LIT>" , None ) <EOL> self . _trace_stmt = kwargs . pop ( "<STR_LIT>" , None ) <EOL> _cursor . __init__ ( self , * args , ** kwargs ) <EOL> self . _tracer = self . _tracer or getattr ( self . connection , "<STR_LIT>" , None ) <EOL> self . _trace_stmt = self . _trace_stmt or getattr ( self . connection , "<STR_LIT>" , None ) <EOL> def execute ( self , query , vars = None ) : <EOL> query_time = <NUM_LIT:0> <EOL> try : <EOL> t_begin = time . time ( ) <EOL> res = super ( TracingCursor , self ) . execute ( query , vars ) <EOL> query_time = time . time ( ) - t_begin <EOL> return res <EOL> finally : <EOL> if self . _tracer : <EOL> self . _log_call ( self . _tracer , trace_stmt = self . _trace_stmt , query_time = query_time ) <EOL> def callproc ( self , procname , vars = None ) : <EOL> query_time = <NUM_LIT:0> <EOL> try : <EOL> t_begin = time . time ( ) <EOL> res = super ( TracingCursor , self ) . callproc ( procname , vars ) <EOL> query_time = time . time ( ) - t_begin <EOL> return res <EOL> finally : <EOL> if self . _tracer : <EOL> self . _log_call ( self . _tracer , trace_stmt = self . _trace_stmt , query_time = query_time ) <EOL> def fetchall ( self ) : <EOL> query_time = <NUM_LIT:0> <EOL> try : <EOL> t_begin = time . time ( ) <EOL> res = super ( TracingCursor , self ) . fetchall ( ) <EOL> query_time = time . time ( ) - t_begin <EOL> return res <EOL> finally : <EOL> log_entry = getattr ( self , "<STR_LIT>" , None ) <EOL> if log_entry and log_entry . get ( "<STR_LIT>" , "<STR_LIT>" ) == self . query and "<STR_LIT>" in log_entry : <EOL> log_entry [ "<STR_LIT>" ] += query_time <EOL> def _log_call ( self , tracer , trace_stmt = None , query_time = None ) : <EOL> statement = trace_stmt or self . query <EOL> status = self . rowcount <EOL> log_entry = dict ( <EOL> statement = statement , <EOL> status = status , <EOL> ) <EOL> if query_time is not None : <EOL> log_entry [ "<STR_LIT>" ] = query_time <EOL> tracer . log_call ( log_entry , include_stack = True ) <EOL> self . _current_entry = log_entry <EOL> return log_entry <EOL> class StatementBuilder ( object ) : <EOL> def __init__ ( self ) : <EOL> self . statement = None <EOL> self . st_frag = [ ] <EOL> self . statement_args = { } <EOL> def append ( self , * fragments ) : <EOL> self . st_frag . extend ( fragments ) <EOL> def append_value ( self , fragment , valkey , value , count ) : <EOL> if count : <EOL> self . st_frag . append ( "<STR_LIT:U+002C>" ) <EOL> self . st_frag . append ( fragment ) <EOL> self . statement_args [ valkey + str ( count ) ] = value <EOL> def build ( self ) : <EOL> self . statement = "<STR_LIT>" . join ( self . st_frag ) <EOL> return self . statement , self . statement_args </s>
