From ac69b228b058b4a8cf74f5c5a6b56d4f4a9a0b0d Mon Sep 17 00:00:00 2001
From: Ann Kamyshnikova <akamyshnikova@mirantis.com>
Date: Mon, 11 Jul 2016 14:43:35 +0300
Subject: [PATCH] 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                 | 25 ++++++++++---------------
 neutron/tests/unit/db/test_api.py | 17 +++++++++++++++++
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/neutron/db/api.py b/neutron/db/api.py
index d85fdcbae5a..5e3ed915ee1 100644
--- a/neutron/db/api.py
+++ b/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)
diff --git a/neutron/tests/unit/db/test_api.py b/neutron/tests/unit/db/test_api.py
index f22e53491e2..a84c1b6f0a9 100644
--- a/neutron/tests/unit/db/test_api.py
+++ b/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')