From a6ad8da299e64beb534ed6eb88c907a4a8ad96e4 Mon Sep 17 00:00:00 2001
From: Boris Pavlovic <boris@pavlovic.me>
Date: Wed, 9 Jul 2014 17:55:53 +0400
Subject: [PATCH] Improve tracing of sqlalchemy

Use before_cursor_execute and after_cursor_execute instead of
before_execute and after_execute.

1) Migration with custom types will work now (even if tracing is enabled)
2) We don't need to render 2 times SQL expressesion => less overhead

Change-Id: I7c6b6909ce0f15a69ce5caad544e1351a647b472
---
 osprofiler/sqlalchemy.py | 27 +++++++++++----------------
 tests/test_sqlalchemy.py | 27 +++++++++++++--------------
 2 files changed, 24 insertions(+), 30 deletions(-)

diff --git a/osprofiler/sqlalchemy.py b/osprofiler/sqlalchemy.py
index f7c4767..4120a1f 100644
--- a/osprofiler/sqlalchemy.py
+++ b/osprofiler/sqlalchemy.py
@@ -20,11 +20,7 @@ _DISABLED = False
 
 
 def disable():
-    """add_tracing does not add event listeners for sqlalchemy.
-
-    This is quite important in case of sql migrations. Because it's not allowed
-    to add these events.
-    """
+    """Disable tracing of all DB queries. Reduce a lot size of profiles."""
     global _DISABLED
     _DISABLED = True
 
@@ -40,27 +36,26 @@ def add_tracing(sqlalchemy, engine, name):
     """Add tracing to all sqlalchemy calls."""
 
     if not _DISABLED:
-        sqlalchemy.event.listen(engine, 'before_execute',
-                                _before_execute(name))
-        sqlalchemy.event.listen(engine, 'after_execute', _after_execute())
+        sqlalchemy.event.listen(engine, 'before_cursor_execute',
+                                _before_cursor_execute(name))
+        sqlalchemy.event.listen(engine, 'after_cursor_execute',
+                                _after_cursor_execute())
 
 
-def _before_execute(name):
+def _before_cursor_execute(name):
     """Add listener that will send trace info before query is executed."""
 
-    def handler(conn, clauseelement, multiparams, params):
-        info = {"db.statement": str(clauseelement),
-                "db.multiparams": str(multiparams),
-                "db.params": str(params)}
+    def handler(conn, cursor, statement, params, context, executemany):
+        info = {"db.statement": statement, "db.params": params}
         profiler.start(name, info=info)
 
     return handler
 
 
-def _after_execute():
+def _after_cursor_execute():
     """Add listener that will send trace info after query is executed."""
 
-    def handler(conn, clauseelement, multiparams, params, result):
-        profiler.stop(info=None)
+    def handler(conn, cursor, statement, params, context, executemany):
+        profiler.stop()
 
     return handler
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
index c53042e..7c56049 100644
--- a/tests/test_sqlalchemy.py
+++ b/tests/test_sqlalchemy.py
@@ -24,24 +24,23 @@ class SqlalchemyTracingTestCase(test.TestCase):
 
     @mock.patch("osprofiler.sqlalchemy.profiler")
     def test_before_execute(self, mock_profiler):
-        handler = sqlalchemy._before_execute("sql")
+        handler = sqlalchemy._before_cursor_execute("sql")
 
-        handler(mock.MagicMock(), 1, 2, 3)
+        handler(mock.MagicMock(), 1, 2, 3, 4, 5)
         expected_info = {
-            "db.statement": "1",
-            "db.multiparams": "2",
-            "db.params": "3"
+            "db.statement": 2,
+            "db.params": 3
         }
         mock_profiler.start.assert_called_once_with("sql", info=expected_info)
 
     @mock.patch("osprofiler.sqlalchemy.profiler")
     def test_after_execute(self, mock_profiler):
-        handler = sqlalchemy._after_execute()
-        handler(mock.MagicMock(), 1, 2, 3, mock.MagicMock())
-        mock_profiler.stop.assert_called_once_with(info=None)
+        handler = sqlalchemy._after_cursor_execute()
+        handler(mock.MagicMock(), 1, 2, 3, 4, 5)
+        mock_profiler.stop.assert_called_once_with()
 
-    @mock.patch("osprofiler.sqlalchemy._before_execute")
-    @mock.patch("osprofiler.sqlalchemy._after_execute")
+    @mock.patch("osprofiler.sqlalchemy._before_cursor_execute")
+    @mock.patch("osprofiler.sqlalchemy._after_cursor_execute")
     def test_add_tracing(self, mock_after_exc, mock_before_exc):
         sa = mock.MagicMock()
         engine = mock.MagicMock()
@@ -54,13 +53,13 @@ class SqlalchemyTracingTestCase(test.TestCase):
         mock_before_exc.assert_called_once_with("sql")
         mock_after_exc.assert_called_once_with()
         expected_calls = [
-            mock.call(engine, "before_execute", "before"),
-            mock.call(engine, "after_execute", "after")
+            mock.call(engine, "before_cursor_execute", "before"),
+            mock.call(engine, "after_cursor_execute", "after")
         ]
         self.assertEqual(sa.event.listen.call_args_list, expected_calls)
 
-    @mock.patch("osprofiler.sqlalchemy._before_execute")
-    @mock.patch("osprofiler.sqlalchemy._after_execute")
+    @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):
         sqlalchemy.disable()