pgtoolkit/pgmanager.py
changeset 33 bd0beda49bcb
parent 31 c2e6e24b83d9
child 36 e67101c22e83
equal deleted inserted replaced
32:d59c473c9ad7 33:bd0beda49bcb
   117 
   117 
   118     def execute(self, query, args=None):
   118     def execute(self, query, args=None):
   119         try:
   119         try:
   120             return super(Cursor, self).execute(query, args)
   120             return super(Cursor, self).execute(query, args)
   121         finally:
   121         finally:
   122             if self.query:
   122             self._log_query()
   123                 log_sql.info(self.query.decode('utf8'))
       
   124 
   123 
   125     def callproc(self, procname, args=None):
   124     def callproc(self, procname, args=None):
   126         try:
   125         try:
   127             return super(Cursor, self).callproc(procname, args)
   126             return super(Cursor, self).callproc(procname, args)
   128         finally:
   127         finally:
   129             if self.query:
   128             self._log_query()
   130                 log_sql.info(self.query.decode('utf8'))
       
   131 
   129 
   132     def row_dict(self, row, lstrip=None):
   130     def row_dict(self, row, lstrip=None):
   133         adjustname = lambda a: a
   131         adjustname = lambda a: a
   134         if lstrip:
   132         if lstrip:
   135             adjustname = lambda a: a.lstrip(lstrip)
   133             adjustname = lambda a: a.lstrip(lstrip)
   162     def fetchall_adapted(self, lstrip=None):
   160     def fetchall_adapted(self, lstrip=None):
   163         '''Like fetchall_dict() but values are quoted for direct inclusion in SQL query.'''
   161         '''Like fetchall_dict() but values are quoted for direct inclusion in SQL query.'''
   164         rows = super(Cursor, self).fetchall()
   162         rows = super(Cursor, self).fetchall()
   165         return [self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) for row in rows]
   163         return [self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) for row in rows]
   166 
   164 
       
   165     def _log_query(self):
       
   166         if self.query:
       
   167             name = self.connection.name if hasattr(self.connection, 'name') else '-'
       
   168             log_sql.info('[%s] %s' % (name, self.query.decode('utf8')))
   167 
   169 
   168 class Connection(psycopg2.extensions.connection):
   170 class Connection(psycopg2.extensions.connection):
   169 
   171 
   170     def cursor(self, name=None):
   172     def cursor(self, name=None):
   171         if name is None:
   173         if name is None:
   264                     conn = None
   266                     conn = None
   265 
   267 
   266             if conn is None:
   268             if conn is None:
   267                 ci = self.conn_known[name]
   269                 ci = self.conn_known[name]
   268                 conn = self._connect(ci)
   270                 conn = self._connect(ci)
       
   271                 # add our name to connection instance (this is then logged with SQL queries)
       
   272                 conn.name = name
   269         finally:
   273         finally:
   270             self.lock.release()
   274             self.lock.release()
   271         return conn
   275         return conn
   272 
   276 
   273     def put_conn(self, conn, name='default'):
   277     def put_conn(self, conn, name='default'):
   290                 conn.close()
   294                 conn.close()
   291                 return
   295                 return
   292 
   296 
   293             # connection returned to the pool must not be in transaction
   297             # connection returned to the pool must not be in transaction
   294             if conn.get_transaction_status() != psycopg2.extensions.TRANSACTION_STATUS_IDLE:
   298             if conn.get_transaction_status() != psycopg2.extensions.TRANSACTION_STATUS_IDLE:
   295                 conn.rollback()
   299                 try:
       
   300                     conn.rollback()
       
   301                 except OperationalError:
       
   302                     if not conn.closed:
       
   303                         conn.close()
       
   304                     return
   296 
   305 
   297             self.conn_pool[name].append(conn)
   306             self.conn_pool[name].append(conn)
   298         finally:
   307         finally:
   299             self.lock.release()
   308             self.lock.release()
   300 
   309