Refactor setting OSprofiler for db calls

This change is depends on change in oslo.db that adds option to
set hook for osprofiler.

Depends-On: I78bef4979c2000d05658ce17d0348cd0a10c24d9

Related-bug: #1600739

Related-bug: #1520719

Change-Id: Ie971654f988c98015edc7a59cf00b27e83c0c1b7

neutron/db/api.py

@@ -29,10 +29,17 @@ from sqlalchemy.orm import exc
 from neutron.common import exceptions
 from neutron.common import profiler  # noqa
 
-context_manager = enginefacade.transaction_context()
-context_manager.configure(sqlite_fk=True)
 
-_PROFILER_INITIALIZED = False
+def set_hook(engine):
+    if cfg.CONF.profiler.enabled and cfg.CONF.profiler.trace_sqlalchemy:
+        osprofiler.sqlalchemy.add_tracing(sqlalchemy, engine, 'neutron.db')
+
+
+context_manager = enginefacade.transaction_context()
+
+context_manager.configure(sqlite_fk=True)
+context_manager.append_on_engine_create(set_hook)
+
 
 MAX_RETRIES = 10
 
@@ -77,16 +84,6 @@ def _is_nested_instance(e, etypes):
             any(_is_nested_instance(i, etypes) for i in e.inner_exceptions))
 
 
-# TODO(akamyshnikova) this code should be in oslo.db
-def _set_profiler():
-    global _PROFILER_INITIALIZED
-    if (not _PROFILER_INITIALIZED and cfg.CONF.profiler.enabled and
-            cfg.CONF.profiler.trace_sqlalchemy):
-        _PROFILER_INITIALIZED = True
-        osprofiler.sqlalchemy.add_tracing(
-            sqlalchemy, context_manager.get_legacy_facade().get_engine(), "db")
-
-
 @contextlib.contextmanager
 def exc_to_retry(exceptions):
     try:
@@ -102,7 +99,6 @@ def exc_to_retry(exceptions):
 # connections will be updated, this won't be needed
 def get_engine():
     """Helper method to grab engine."""
-    _set_profiler()
     return context_manager.get_legacy_facade().get_engine()
 
 
@@ -117,7 +113,6 @@ def dispose():
 # connections will be updated, this won't be needed
 def get_session(autocommit=True, expire_on_commit=False, use_slave=False):
     """Helper method to grab session."""
-    _set_profiler()
     return context_manager.get_legacy_facade().get_session(
         autocommit=autocommit, expire_on_commit=expire_on_commit,
         use_slave=use_slave)

neutron/tests/unit/db/test_api.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
+from oslo_config import cfg
 from oslo_db import exception as db_exc
+import osprofiler
+import sqlalchemy
 from sqlalchemy.orm import exc
 import testtools
 
@@ -92,3 +96,16 @@ class TestDeadLockDecorator(base.BaseTestCase):
         e = db_exc.DBError("(pymysql.err.InternalError) (1305, u'SAVEPOINT "
                            "sa_savepoint_1 does not exist')")
         self.assertIsNone(self._decorated_function(1, e))
+
+
+class TestCommonDBfunctions(base.BaseTestCase):
+
+    def test_set_hook(self):
+        with mock.patch.object(osprofiler.sqlalchemy,
+                               'add_tracing') as profiler:
+            cfg.CONF.set_override('enabled', True, group='profiler')
+            cfg.CONF.set_override('trace_sqlalchemy', True, group='profiler')
+            engine_mock = mock.Mock()
+            db_api.set_hook(engine_mock)
+            profiler.assert_called_with(sqlalchemy, engine_mock,
+                                        'neutron.db')