pgtoolkit/pgmanager.py
changeset 33 bd0beda49bcb
parent 31 c2e6e24b83d9
child 36 e67101c22e83
--- a/pgtoolkit/pgmanager.py	Fri Mar 23 14:54:04 2012 +0100
+++ b/pgtoolkit/pgmanager.py	Wed Mar 28 17:25:18 2012 +0200
@@ -119,15 +119,13 @@
         try:
             return super(Cursor, self).execute(query, args)
         finally:
-            if self.query:
-                log_sql.info(self.query.decode('utf8'))
+            self._log_query()
 
     def callproc(self, procname, args=None):
         try:
             return super(Cursor, self).callproc(procname, args)
         finally:
-            if self.query:
-                log_sql.info(self.query.decode('utf8'))
+            self._log_query()
 
     def row_dict(self, row, lstrip=None):
         adjustname = lambda a: a
@@ -164,6 +162,10 @@
         rows = super(Cursor, self).fetchall()
         return [self.row_dict([self.mogrify('%s', [x]).decode('utf8') for x in row], lstrip) for row in rows]
 
+    def _log_query(self):
+        if self.query:
+            name = self.connection.name if hasattr(self.connection, 'name') else '-'
+            log_sql.info('[%s] %s' % (name, self.query.decode('utf8')))
 
 class Connection(psycopg2.extensions.connection):
 
@@ -266,6 +268,8 @@
             if conn is None:
                 ci = self.conn_known[name]
                 conn = self._connect(ci)
+                # add our name to connection instance (this is then logged with SQL queries)
+                conn.name = name
         finally:
             self.lock.release()
         return conn
@@ -292,7 +296,12 @@
 
             # connection returned to the pool must not be in transaction
             if conn.get_transaction_status() != psycopg2.extensions.TRANSACTION_STATUS_IDLE:
-                conn.rollback()
+                try:
+                    conn.rollback()
+                except OperationalError:
+                    if not conn.closed:
+                        conn.close()
+                    return
 
             self.conn_pool[name].append(conn)
         finally: