Merge "Add a hook to process newly created engines"

2 files changed, 24 insertions, 1 deletions

diff --git a/oslo_db/sqlalchemy/enginefacade.py b/oslo_db/sqlalchemy/enginefacade.py
index cd6745b..9929304 100644
--- a/oslo_db/sqlalchemy/enginefacade.py
+++ b/oslo_db/sqlalchemy/enginefacade.py
@@ -154,7 +154,8 @@ class _TransactionFactory(object):
             'rollback_reader_sessions': False,
             }
         self._facade_cfg = {
-            'synchronous_reader': True
+            'synchronous_reader': True,
+            'on_engine_create': [],
         }
 
         # other options that are defined in oslo.db.options.database_opts
@@ -373,6 +374,8 @@ class _TransactionFactory(object):
                 "No sql_connection parameter is established")
         engine = engines.create_engine(
             sql_connection=sql_connection, **engine_kwargs)
+        for hook in self._facade_cfg['on_engine_create']:
+            hook(engine)
         sessionmaker = orm.get_maker(engine=engine, **maker_kwargs)
         return engine, sessionmaker
 
@@ -637,6 +640,10 @@ class _TransactionContextManager(object):
         """
         self._factory.configure(**kw)
 
+    def append_on_engine_create(self, fn):
+        """Append a listener function to _facade_cfg["on_engine_create"]"""
+        self._factory._facade_cfg['on_engine_create'].append(fn)
+
     def get_legacy_facade(self):
         """Return a :class:`.LegacyEngineFacade` for factory from this context.
 
diff --git a/oslo_db/tests/sqlalchemy/test_enginefacade.py b/oslo_db/tests/sqlalchemy/test_enginefacade.py
index 1365f1c..72663d2 100644
--- a/oslo_db/tests/sqlalchemy/test_enginefacade.py
+++ b/oslo_db/tests/sqlalchemy/test_enginefacade.py
@@ -1883,5 +1883,21 @@ class ConfigOptionsTest(oslo_test_base.BaseTestCase):
         )
 
 
+class TestTransactionFactoryCallback(oslo_test_base.BaseTestCase):
+
+    def test_setup_for_connection_called_with_profiler(self):
+        context_manager = enginefacade.transaction_context()
+        context_manager.configure(connection='sqlite://')
+        hook = mock.Mock()
+        context_manager.append_on_engine_create(hook)
+        self.assertEqual(
+            [hook], context_manager._factory._facade_cfg['on_engine_create'])
+
+        @context_manager.reader
+        def go(context):
+            hook.assert_called_once_with(context.session.bind)
+
+        go(oslo_context.RequestContext())
+
 # TODO(zzzeek): test configuration options, e.g. like
 # test_sqlalchemy->test_creation_from_config