Issue 120 add pool close

happybase/pool.py

@@ -90,6 +90,33 @@ def _return_connection(self, connection):
         """Return a connection to the pool."""
         self._queue.put(connection)
 
+    def close_connections(self, timeout=None):
+        """
+        Attempts to politely close all connections in the pool.
+        Waits for used connections to become available or until the
+        timeout is exceeded.
+
+        :param int timeout: number of seconds to wait for a connection to
+                            become available (optional)
+        :return: None
+        """
+        if timeout:
+            if not isinstance(timeout, int):
+                raise TypeError("close_connections 'timeout' arg must be an integer")
+
+            if not timeout > 0:
+                raise ValueError("close_connections 'timeout' arg must be greater than zero")
+
+        while not self._queue.empty():
+            try:
+                conn = self._queue.get(block=True, timeout=timeout)
+                conn.close()
+                conn = None
+            except Queue.Empty:
+                raise NoConnectionsAvailable(
+                    "Closing connections failed: No connection available from "
+                    "pool within specified timeout of {}".format(timeout))
+
     @contextlib.contextmanager
     def connection(self, timeout=None):
         """

tests/test_api.py

@@ -556,6 +556,17 @@ def run():
         t.start()
         t.join()
 
+def test_pool_close():
+    pool = ConnectionPool(size=3, **connection_kwargs)
+
+    with pool.connection():
+        with assert_raises(TypeError):
+            pool.close_connections(timeout='foo')
+
+        with assert_raises(ValueError):
+            pool.close_connections(timeout=-1)
+
+        pool.close_connections()
 
 if __name__ == '__main__':
     import logging