Add a hook to process newly created engines

Add on_engine_create argument to transaction_context to enable users
to take some actions with newly created engines. Main intended usage
is providing oslo.db and osprofiler users with a way to hook them
together without making them depend on each other.

Closes-bug: #1600739

Change-Id: I78bef4979c2000d05658ce17d0348cd0a10c24d9

2 files changed, 24 insertions, 1 deletions

diff --git a/oslo_db/sqlalchemy/enginefacade.py b/oslo_db/sqlalchemy/enginefacade.py
index c935602..79dbc56 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
@@ -363,6 +364,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
 
@@ -627,6 +630,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 d27f7bf..070558d 100644
--- a/oslo_db/tests/sqlalchemy/test_enginefacade.py
+++ b/oslo_db/tests/sqlalchemy/test_enginefacade.py
@@ -1840,5 +1840,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