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 is not None:
+ 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=0)
+
+ pool.close_connections()
 
 if __name__ == '__main__':
  import logging