Merge 4e6ab07b30 into 90ee1ebba5

lib/pool.py

@@ -32,6 +32,20 @@ class PoolError(psycopg2.Error):
     pass
 
 
+class ConnectionContext(object):
+
+    def __init__(self, pool):
+        self.pool = pool
+        self._conn = None
+
+    def __enter__(self):
+        self._conn = self.pool.getconn()
+        return self._conn
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.pool.putconn(self._conn)
+
+
 class AbstractConnectionPool(object):
     """Generic key-based pooling code."""
 
@@ -57,6 +71,9 @@ class AbstractConnectionPool(object):
         for i in range(self.minconn):
             self._connect()
 
+    def __call__(self):
+        return ConnectionContext(self)
+
     def _connect(self, key=None):
         """Create a new connection and assign it to 'key' if not None."""
         conn = psycopg2.connect(*self._args, **self._kwargs)