pgtoolkit/pgmanager.py
changeset 24 5664afa530e5
parent 23 dc2dbe872fc8
child 26 7f219da7ab71
--- a/pgtoolkit/pgmanager.py	Wed Dec 14 16:29:33 2011 +0100
+++ b/pgtoolkit/pgmanager.py	Thu Dec 15 18:21:41 2011 +0100
@@ -72,6 +72,7 @@
 from contextlib import contextmanager
 import logging
 import threading
+import multiprocessing
 import select
 import socket
 
@@ -113,13 +114,15 @@
         try:
             return super(Cursor, self).execute(query, args)
         finally:
-            log.debug(self.query.decode('utf8'))
+            if self.query:
+                log.debug(self.query.decode('utf8'))
 
     def callproc(self, procname, args=None):
         try:
             return super(Cursor, self).callproc(procname, args)
         finally:
-            log.debug(self.query.decode('utf8'))
+            if self.query:
+                log.debug(self.query.decode('utf8'))
 
     def row_dict(self, row, lstrip=None):
         adjustname = lambda a: a
@@ -182,8 +185,9 @@
 
     def __init__(self):
         self.conn_known = {}  # available connections
-        self.conn_pool = {}
-        self.lock = threading.Lock()
+        self.conn_pool = {}  # active connetions
+        self.lock = threading.Lock()  # mutual exclusion for threads
+        self.pid = multiprocessing.current_process().pid  # forking check
 
     def __del__(self):
         for conn in tuple(self.conn_known.keys()):
@@ -240,6 +244,7 @@
 
     def get_conn(self, name='default'):
         '''Get connection of name 'name' from pool.'''
+        self._check_fork()
         self.lock.acquire()
         try:
             if not name in self.conn_known:
@@ -355,6 +360,26 @@
                 return psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE
             raise PgManagerError('Unknown isolation level name: "%s"', level)
         return level
+    
+    def _check_fork(self):
+        '''Check if process was forked (PID has changed).
+        
+        If it was, clean parent's connections.
+        New connections are created for children.
+        Known connection credentials are inherited, but not shared.
+        
+        '''
+        if self.pid == multiprocessing.current_process().pid:
+            # PID has not changed
+            return
+        
+        # update saved PID
+        self.pid = multiprocessing.current_process().pid
+        # reinitialize lock
+        self.lock = threading.Lock()
+        # clean parent's connections
+        for name in self.conn_pool:
+            self.conn_pool[name] = []
 
     @classmethod
     def get_instance(cls):