Add hepler to trace sessions
Lots of OpenStack projects are moving from direct usage of engine to enginefacade module usage with its _TransactionContextManager. Sessions for the context are created via .using() method of this manager. Change-Id: Ie517524a10713e09bb2491f2d76b651b20ecd4b6
This commit is contained in:
parent
8be653361f
commit
eab2574110
@ -13,6 +13,8 @@
|
||||
# License for the specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import contextlib
|
||||
|
||||
from osprofiler import profiler
|
||||
|
||||
|
||||
@ -42,6 +44,15 @@ def add_tracing(sqlalchemy, engine, name):
|
||||
_after_cursor_execute())
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def wrap_session(sqlalchemy, sess):
|
||||
with sess as s:
|
||||
if not getattr(s.bind, "traced", False):
|
||||
add_tracing(sqlalchemy, s.bind, "db")
|
||||
s.bind.traced = True
|
||||
yield s
|
||||
|
||||
|
||||
def _before_cursor_execute(name):
|
||||
"""Add listener that will send trace info before query is executed."""
|
||||
|
||||
|
@ -13,6 +13,7 @@
|
||||
# License for the specific language governing permissions and limitations
|
||||
# under the License.
|
||||
|
||||
import contextlib
|
||||
import mock
|
||||
|
||||
from osprofiler import sqlalchemy
|
||||
@ -54,6 +55,36 @@ class SqlalchemyTracingTestCase(test.TestCase):
|
||||
]
|
||||
self.assertEqual(sa.event.listen.call_args_list, expected_calls)
|
||||
|
||||
@mock.patch("osprofiler.sqlalchemy._before_cursor_execute")
|
||||
@mock.patch("osprofiler.sqlalchemy._after_cursor_execute")
|
||||
def test_wrap_session(self, mock_after_exc, mock_before_exc):
|
||||
sa = mock.MagicMock()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _session():
|
||||
session = mock.MagicMock()
|
||||
# current engine object stored within the session
|
||||
session.bind = mock.MagicMock()
|
||||
session.bind.traced = None
|
||||
yield session
|
||||
|
||||
mock_before_exc.return_value = "before"
|
||||
mock_after_exc.return_value = "after"
|
||||
|
||||
session = sqlalchemy.wrap_session(sa, _session())
|
||||
|
||||
with session as sess:
|
||||
pass
|
||||
|
||||
mock_before_exc.assert_called_once_with("db")
|
||||
mock_after_exc.assert_called_once_with()
|
||||
expected_calls = [
|
||||
mock.call(sess.bind, "before_cursor_execute", "before"),
|
||||
mock.call(sess.bind, "after_cursor_execute", "after")
|
||||
]
|
||||
|
||||
self.assertEqual(sa.event.listen.call_args_list, expected_calls)
|
||||
|
||||
@mock.patch("osprofiler.sqlalchemy._before_cursor_execute")
|
||||
@mock.patch("osprofiler.sqlalchemy._after_cursor_execute")
|
||||
def test_disable_and_enable(self, mock_after_exc, mock_before_exc):
|
||||
|
Loading…
Reference in New Issue
Block a user