Merge "Use SQLAlchemy cursor execute events for tracing"

2 files changed, 57 insertions, 40 deletions

diff --git a/oslo/db/sqlalchemy/session.py b/oslo/db/sqlalchemy/session.py
index 07a062e..ba6b8f7 100644
--- a/oslo/db/sqlalchemy/session.py
+++ b/oslo/db/sqlalchemy/session.py
@@ -470,8 +470,8 @@ def create_engine(sql_connection, sqlite_fk=False, mysql_sql_mode=None,
                                     _synchronous_switch_listener)
         sqlalchemy.event.listen(engine, 'connect', _add_regexp_listener)
 
-    if connection_trace and engine.dialect.dbapi.__name__ == 'MySQLdb':
-        _patch_mysqldb_with_stacktrace_comments()
+    if connection_trace:
+        _add_trace_comments(engine)
 
     # register alternate exception handler
     exc_filters.register_engine(engine)
@@ -525,47 +525,48 @@ def get_maker(engine, autocommit=True, expire_on_commit=False):
                                        query_cls=Query)
 
 
-def _patch_mysqldb_with_stacktrace_comments():
-    """Adds current stack trace as a comment in queries.
+def _add_trace_comments(engine):
+    """Augment statements with a trace of the immediate calling code
+    for a given statement.
 
-    Patches MySQLdb.cursors.BaseCursor._do_query.
     """
-    import MySQLdb.cursors
-    import traceback
-
-    old_mysql_do_query = MySQLdb.cursors.BaseCursor._do_query
-
-    def _do_query(self, q):
-        stack = ''
-        for filename, line, method, function in traceback.extract_stack():
-            # exclude various common things from trace
-            if filename.endswith('session.py') and method == '_do_query':
-                continue
-            if filename.endswith('api.py') and method == 'wrapper':
-                continue
-            if filename.endswith('utils.py') and method == '_inner':
-                continue
-            if filename.endswith('exception.py') and method == '_wrap':
-                continue
-            # db/api is just a wrapper around db/sqlalchemy/api
-            if filename.endswith('db/api.py'):
-                continue
-            # only trace inside oslo
-            index = filename.rfind('oslo')
-            if index == -1:
-                continue
-            stack += "File:%s:%s Method:%s() Line:%s | " \
-                     % (filename[index:], line, method, function)
-
-        # strip trailing " | " from stack
-        if stack:
-            stack = stack[:-3]
-            qq = "%s /* %s */" % (q, stack)
-        else:
-            qq = q
-        old_mysql_do_query(self, qq)
 
-    setattr(MySQLdb.cursors.BaseCursor, '_do_query', _do_query)
+    import os
+    import sys
+    import traceback
+    target_paths = set([
+        os.path.dirname(sys.modules['oslo.db'].__file__),
+        os.path.dirname(sys.modules['sqlalchemy'].__file__)
+    ])
+
+    @sqlalchemy.event.listens_for(engine, "before_cursor_execute", retval=True)
+    def before_cursor_execute(conn, cursor, statement, parameters, context,
+            executemany):
+
+        # NOTE(zzzeek) - if different steps per DB dialect are desirable
+        # here, switch out on engine.name for now.
+        stack = traceback.extract_stack()
+        our_line = None
+        for idx, (filename, line, method, function) in enumerate(stack):
+            for tgt in target_paths:
+                if filename.startswith(tgt):
+                    our_line = idx
+                    break
+            if our_line:
+                break
+
+        if our_line:
+            trace = "; ".join(
+                "File: %s (%s) %s" % (
+                    line[0], line[1], line[2]
+                )
+                # include three lines of context.
+                for line in stack[our_line - 3:our_line]
+
+            )
+            statement = "%s  -- %s" % (statement, trace)
+
+        return statement, parameters
 
 
 class EngineFacade(object):
diff --git a/tests/sqlalchemy/test_sqlalchemy.py b/tests/sqlalchemy/test_sqlalchemy.py
index fac17ae..a6278d5 100644
--- a/tests/sqlalchemy/test_sqlalchemy.py
+++ b/tests/sqlalchemy/test_sqlalchemy.py
@@ -451,3 +451,19 @@ class MysqlSetCallbackTest(oslo_test.BaseTestCase):
             "SELECT * FROM bar",
         ]
         self.assertEqual(exp_calls, engine._execs)
+
+
+class PatchStacktraceTest(test_base.DbTestCase):
+
+    def test_trace(self):
+        engine = self.engine
+        engine.connect()
+        with mock.patch.object(engine.dialect, "do_execute") as mock_exec:
+            session._add_trace_comments(engine)
+
+            engine.execute("select * from table")
+
+        call = mock_exec.mock_calls[0]
+
+        # we're the caller, see that we're in there
+        self.assertTrue("tests/sqlalchemy/test_sqlalchemy.py" in call[1][1])