diff --git a/happybase/pool.py b/happybase/pool.py index aa0da99..5ec4e04 100644 --- a/happybase/pool.py +++ b/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): """ diff --git a/tests/test_api.py b/tests/test_api.py index cbfdbd0..3b6913f 100644 --- a/tests/test_api.py +++ b/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