<s> """<STR_LIT>""" <EOL> import sys <EOL> from psycopg2 import connect <EOL> from psycopg2 . extras import DictConnection <EOL> from . pycompat import PY2 <EOL> if PY2 : <EOL> from psycopg2 . extensions import register_type , UNICODE <EOL> register_type ( UNICODE ) <EOL> class DbConnection ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , dbname , user = None , pswd = None , host = None , port = None ) : <EOL> """<STR_LIT>""" <EOL> self . dbname = dbname <EOL> self . user = '<STR_LIT>' if user is None else "<STR_LIT>" % user <EOL> self . pswd = '<STR_LIT>' if pswd is None else "<STR_LIT>" % pswd <EOL> self . host = '<STR_LIT>' if host is None else "<STR_LIT>" % host <EOL> self . port = '<STR_LIT>' if port is None else "<STR_LIT>" % port <EOL> self . conn = None <EOL> def connect ( self ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> self . conn = connect ( "<STR_LIT>" % ( <EOL> self . host , self . port , self . dbname , self . user , self . pswd ) , <EOL> connection_factory = DictConnection ) <EOL> except Exception as exc : <EOL> if str ( exc ) [ : <NUM_LIT:6> ] == '<STR_LIT>' : <EOL> sys . exit ( "<STR_LIT>" % str ( exc ) [ <NUM_LIT:8> : ] ) <EOL> else : <EOL> raise exc <EOL> def close ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . conn and not self . conn . closed : <EOL> self . conn . close ( ) <EOL> self . conn = None <EOL> def commit ( self ) : <EOL> """<STR_LIT>""" <EOL> self . conn . commit ( ) <EOL> def rollback ( self ) : <EOL> """<STR_LIT>""" <EOL> self . conn . rollback ( ) <EOL> def execute ( self , query , args = None ) : <EOL> """<STR_LIT>""" <EOL> if self . conn is None or self . conn . closed : <EOL> self . connect ( ) <EOL> curs = self . conn . cursor ( ) <EOL> try : <EOL> curs . execute ( query , args ) <EOL> except Exception as exc : <EOL> self . conn . rollback ( ) <EOL> curs . close ( ) <EOL> raise exc <EOL> return curs <EOL> def fetchone ( self , query , args = None ) : <EOL> """<STR_LIT>""" <EOL> curs = self . execute ( query , args ) <EOL> row = curs . fetchone ( ) <EOL> curs . close ( ) <EOL> return row <EOL> def fetchall ( self , query , args = None ) : <EOL> """<STR_LIT>""" <EOL> curs = self . execute ( query , args ) <EOL> rows = curs . fetchall ( ) <EOL> curs . close ( ) <EOL> return rows <EOL> def copy_to ( self , path , table , sep = '<STR_LIT:U+002C>' ) : <EOL> """<STR_LIT>""" <EOL> if self . conn is None or self . conn . closed : <EOL> self . connect ( ) <EOL> with open ( path , '<STR_LIT:w>' ) as f : <EOL> curs = self . conn . cursor ( ) <EOL> try : <EOL> curs . copy_to ( f , table , sep ) <EOL> except : <EOL> curs . close ( ) <EOL> raise <EOL> def sql_copy_to ( self , sql , path ) : <EOL> """<STR_LIT>""" <EOL> if self . conn is None or self . conn . closed : <EOL> self . connect ( ) <EOL> with open ( path , '<STR_LIT:w>' ) as f : <EOL> curs = self . conn . cursor ( ) <EOL> try : <EOL> curs . copy_expert ( sql , f ) <EOL> except : <EOL> curs . close ( ) <EOL> raise <EOL> def copy_from ( self , path , table , sep = '<STR_LIT:U+002C>' ) : <EOL> """<STR_LIT>""" <EOL> if self . conn is None or self . conn . closed : <EOL> self . connect ( ) <EOL> with open ( path , '<STR_LIT:r>' ) as f : <EOL> curs = self . conn . cursor ( ) <EOL> try : <EOL> curs . copy_from ( f , table , sep ) <EOL> except : <EOL> curs . close ( ) <EOL> raise </s>
