From 7b76cdf420289353f5a8e8acb7d49882b35d18ac Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 11 Jul 2016 09:59:15 -0400
Subject: [PATCH] Add dispose_pool() method to enginefacade context, factory

Openstack applications such as Neutron will refer to
the engine.pool.dispose() method in order to remove
any remaining connections, typically as part of
a test suite.  Add a public-facing method to the context
and its underlying factory so this can be called without
underscore access.

Change-Id: I18b16e002a8507eb2eaa3a7f168e853b0ca17aa5
---
 oslo_db/sqlalchemy/enginefacade.py            | 14 ++++++
 oslo_db/tests/sqlalchemy/test_enginefacade.py | 43 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/oslo_db/sqlalchemy/enginefacade.py b/oslo_db/sqlalchemy/enginefacade.py
index c935602..cd6745b 100644
--- a/oslo_db/sqlalchemy/enginefacade.py
+++ b/oslo_db/sqlalchemy/enginefacade.py
@@ -308,6 +308,16 @@ class _TransactionFactory(object):
     def _maker_args_for_conf(self, conf):
         return self._args_for_conf(self._maker_cfg, conf)
 
+    def dispose_pool(self):
+        """Call engine.pool.dispose() on underlying Engine objects."""
+        with self._start_lock:
+            if not self._started:
+                return
+
+            self._writer_engine.pool.dispose()
+            if self._reader_engine is not self._writer_engine:
+                self._reader_engine.pool.dispose()
+
     def _start(self, conf=False, connection=None, slave_connection=None):
         with self._start_lock:
             # self._started has been checked on the outside
@@ -638,6 +648,10 @@ class _TransactionContextManager(object):
 
         return self._factory.get_legacy_facade()
 
+    def dispose_pool(self):
+        """Call engine.pool.dispose() on underlying Engine objects."""
+        self._factory.dispose_pool()
+
     @property
     def replace(self):
         """Modifier to replace the global transaction factory with this one."""
diff --git a/oslo_db/tests/sqlalchemy/test_enginefacade.py b/oslo_db/tests/sqlalchemy/test_enginefacade.py
index d27f7bf..1365f1c 100644
--- a/oslo_db/tests/sqlalchemy/test_enginefacade.py
+++ b/oslo_db/tests/sqlalchemy/test_enginefacade.py
@@ -62,6 +62,7 @@ class SingletonEngine(SingletonOnName):
         super(SingletonEngine, self).__init__(
             "engine",
             connect=mock.Mock(return_value=connection),
+            pool=mock.Mock(),
             url=connection,
             _assert_connection=connection,
             **kw
@@ -417,6 +418,48 @@ class MockFacadeTest(oslo_test_base.BaseTestCase):
                 session.mock_calls,
                 self.sessions.element_for_writer(writer).mock_calls)
 
+    def test_dispose_pool(self):
+        facade = enginefacade.transaction_context()
+
+        facade.configure(
+            connection=self.engine_uri,
+            )
+
+        facade.dispose_pool()
+        self.assertFalse(hasattr(facade._factory, '_writer_engine'))
+
+        facade._factory._start()
+        facade.dispose_pool()
+
+        self.assertEqual(
+            facade._factory._writer_engine.pool.mock_calls,
+            [mock.call.dispose()]
+        )
+
+    def test_dispose_pool_w_reader(self):
+        facade = enginefacade.transaction_context()
+
+        facade.configure(
+            connection=self.engine_uri,
+            slave_connection=self.slave_uri
+        )
+
+        facade.dispose_pool()
+        self.assertFalse(hasattr(facade._factory, '_writer_engine'))
+        self.assertFalse(hasattr(facade._factory, '_reader_engine'))
+
+        facade._factory._start()
+        facade.dispose_pool()
+
+        self.assertEqual(
+            facade._factory._writer_engine.pool.mock_calls,
+            [mock.call.dispose()]
+        )
+        self.assertEqual(
+            facade._factory._reader_engine.pool.mock_calls,
+            [mock.call.dispose()]
+        )
+
     def test_session_reader_decorator(self):
         context = oslo_context.RequestContext()