Add function/sql results to trace info

This patch adds function/sql results to trace info
New arg: hide_result (Boolean True/False)
- True (default): Hide function/sql result by default
- False: Show result in trace info

With database: add_tracing(sqlalchemy, engine, name, hide_result=True)

Demo: https://tovin07.github.io/murano/environment-show-with-results.html

Change-Id: I317dfa04e0109d46c1a5ca1e0a3523cfd8470d78

osprofiler/profiler.py

@@ -87,7 +87,8 @@ def stop(info=None):
         profiler.stop(info=info)
 
 
-def trace(name, info=None, hide_args=False, allow_multiple_trace=True):
+def trace(name, info=None, hide_args=False, hide_result=True,
+          allow_multiple_trace=True):
     """Trace decorator for functions.
 
     Very useful if you would like to add trace point on existing function:
@@ -102,6 +103,9 @@ def trace(name, info=None, hide_args=False, allow_multiple_trace=True):
     :param hide_args: Don't push to trace info args and kwargs. Quite useful
                       if you have some info in args that you wont to share,
                       e.g. passwords.
+    :param hide_result: Boolean value to hide/show function result in trace.
+                        True - hide function result (default).
+                        False - show function result in trace.
     :param allow_multiple_trace: If the wrapped function has already been
                                  traced either allow the new trace to occur
                                  or raise a value error denoting that multiple
@@ -133,24 +137,43 @@ def trace(name, info=None, hide_args=False, allow_multiple_trace=True):
 
         @functools.wraps(f)
         def wrapper(*args, **kwargs):
-            if "name" not in info["function"]:
+            # NOTE(tovin07): Workaround for this issue
+            # F823 local variable 'info'
+            # (defined in enclosing scope on line xxx)
+            # referenced before assignment
+            info_ = info
+            if "name" not in info_["function"]:
                 # Get this once (as it should **not** be changing in
                 # subsequent calls).
-                info["function"]["name"] = reflection.get_callable_name(f)
+                info_["function"]["name"] = reflection.get_callable_name(f)
 
             if not hide_args:
-                info["function"]["args"] = str(args)
+                info_["function"]["args"] = str(args)
-                info["function"]["kwargs"] = str(kwargs)
+                info_["function"]["kwargs"] = str(kwargs)
 
-            with Trace(name, info=info):
+            stop_info = None
-                return f(*args, **kwargs)
+            try:
+                start(name, info=info_)
+                result = f(*args, **kwargs)
+            except Exception as ex:
+                stop_info = {"etype": reflection.get_class_name(ex)}
+                raise
+            else:
+                if not hide_result:
+                    stop_info = {"function": {"result": repr(result)}}
+                return result
+            finally:
+                if stop_info:
+                    stop(info=stop_info)
+                else:
+                    stop()
 
         return wrapper
 
     return decorator
 
 
-def trace_cls(name, info=None, hide_args=False,
+def trace_cls(name, info=None, hide_args=False, hide_result=True,
               trace_private=False, allow_multiple_trace=True,
               trace_class_methods=False, trace_static_methods=False):
     """Trace decorator for instances of class .
@@ -173,6 +196,9 @@ def trace_cls(name, info=None, hide_args=False,
     :param hide_args: Don't push to trace info args and kwargs. Quite useful
                       if you have some info in args that you wont to share,
                       e.g. passwords.
+    :param hide_result: Boolean value to hide/show function result in trace.
+                        True - hide function result (default).
+                        False - show function result in trace.
     :param trace_private: Trace methods that starts with "_". It wont trace
                           methods that starts "__" even if it is turned on.
     :param trace_static_methods: Trace staticmethods. This may be prone to
@@ -226,7 +252,8 @@ def trace_cls(name, info=None, hide_args=False,
             # halfway trace this class).
             _ensure_no_multiple_traced(traceable_attrs)
         for i, (attr_name, attr) in enumerate(traceable_attrs):
-            wrapped_method = trace(name, info=info, hide_args=hide_args)(attr)
+            wrapped_method = trace(name, info=info, hide_args=hide_args,
+                                   hide_result=hide_result)(attr)
             wrapper = traceable_wrappers[i]
             if wrapper is not None:
                 wrapped_method = wrapper(wrapped_method)
@@ -246,6 +273,7 @@ class TracedMeta(type):
     >>>      __trace_args__ = {'name': 'rpc',
     >>>                        'info': None,
     >>>                        'hide_args': False,
+    >>>                        'hide_result': True,
     >>>                        'trace_private': False}
     >>>
     >>>      def my_method(self, some_args):

osprofiler/sqlalchemy.py

@@ -34,14 +34,16 @@ def enable():
     _DISABLED = False
 
 
-def add_tracing(sqlalchemy, engine, name):
+def add_tracing(sqlalchemy, engine, name, hide_result=True):
     """Add tracing to all sqlalchemy calls."""
 
     if not _DISABLED:
         sqlalchemy.event.listen(engine, "before_cursor_execute",
                                 _before_cursor_execute(name))
-        sqlalchemy.event.listen(engine, "after_cursor_execute",
+        sqlalchemy.event.listen(
-                                _after_cursor_execute())
+            engine, "after_cursor_execute",
+            _after_cursor_execute(hide_result=hide_result)
+        )
 
 
 @contextlib.contextmanager
@@ -66,10 +68,24 @@ def _before_cursor_execute(name):
     return handler
 
 
-def _after_cursor_execute():
+def _after_cursor_execute(hide_result=True):
-    """Add listener that will send trace info after query is executed."""
+    """Add listener that will send trace info after query is executed.
+
+    :param hide_result: Boolean value to hide or show SQL result in trace.
+                        True - hide SQL result (default).
+                        False - show SQL result in trace.
+    """
 
     def handler(conn, cursor, statement, params, context, executemany):
-        profiler.stop()
+        if not hide_result:
+            # Add SQL result to trace info in *-stop phase
+            info = {
+                "db": {
+                    "result": str(cursor._rows)
+                }
+            }
+            profiler.stop(info=info)
+        else:
+            profiler.stop()
 
     return handler

osprofiler/tests/unit/test_profiler.py

@@ -185,6 +185,11 @@ def test_fn_exc():
     raise ValueError()
 
 
+@profiler.trace("hide_result", hide_result=False)
+def trace_with_result_func(a, i=10):
+    return (a, i)
+
+
 class TraceDecoratorTestCase(test.TestCase):
 
     @mock.patch("osprofiler.profiler.stop")
@@ -242,6 +247,27 @@ class TraceDecoratorTestCase(test.TestCase):
         mock_start.assert_called_once_with("foo", info=expected_info)
         mock_stop.assert_called_once_with(info=expected_stop_info)
 
+    @mock.patch("osprofiler.profiler.stop")
+    @mock.patch("osprofiler.profiler.start")
+    def test_with_result(self, mock_start, mock_stop):
+        self.assertEqual((1, 2), trace_with_result_func(1, i=2))
+        start_info = {
+            "function": {
+                "name": "osprofiler.tests.unit.test_profiler"
+                        ".trace_with_result_func",
+                "args": str((1,)),
+                "kwargs": str({"i": 2})
+            }
+        }
+
+        stop_info = {
+            "function": {
+                "result": str((1, 2))
+            }
+        }
+        mock_start.assert_called_once_with("hide_result", info=start_info)
+        mock_stop.assert_called_once_with(info=stop_info)
+
 
 class FakeTracedCls(object):
 

osprofiler/tests/unit/test_sqlalchemy.py

@@ -37,6 +37,19 @@ class SqlalchemyTracingTestCase(test.TestCase):
         handler(mock.MagicMock(), 1, 2, 3, 4, 5)
         mock_profiler.stop.assert_called_once_with()
 
+    @mock.patch("osprofiler.sqlalchemy.profiler")
+    def test_after_execute_with_sql_result(self, mock_profiler):
+        handler = sqlalchemy._after_cursor_execute(hide_result=False)
+        cursor = mock.MagicMock()
+        cursor._rows = (1,)
+        handler(1, cursor, 2, 3, 4, 5)
+        info = {
+            "db": {
+                "result": str(cursor._rows)
+            }
+        }
+        mock_profiler.stop.assert_called_once_with(info=info)
+
     @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):
@@ -49,7 +62,8 @@ class SqlalchemyTracingTestCase(test.TestCase):
         sqlalchemy.add_tracing(sa, engine, "sql")
 
         mock_before_exc.assert_called_once_with("sql")
-        mock_after_exc.assert_called_once_with()
+        # Default set hide_result=True
+        mock_after_exc.assert_called_once_with(hide_result=True)
         expected_calls = [
             mock.call(engine, "before_cursor_execute", "before"),
             mock.call(engine, "after_cursor_execute", "after")
@@ -78,7 +92,8 @@ class SqlalchemyTracingTestCase(test.TestCase):
             pass
 
         mock_before_exc.assert_called_once_with("db")
-        mock_after_exc.assert_called_once_with()
+        # Default set hide_result=True
+        mock_after_exc.assert_called_once_with(hide_result=True)
         expected_calls = [
             mock.call(sess.bind, "before_cursor_execute", "before"),
             mock.call(sess.bind, "after_cursor_execute", "after")
@@ -86,6 +101,28 @@ 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")
+    @mock.patch("osprofiler.profiler")
+    def test_with_sql_result(self, mock_profiler,
+                             mock_after_exc, mock_before_exc):
+        sa = mock.MagicMock()
+        engine = mock.MagicMock()
+
+        mock_before_exc.return_value = "before"
+        mock_after_exc.return_value = "after"
+
+        sqlalchemy.add_tracing(sa, engine, "sql", hide_result=False)
+
+        mock_before_exc.assert_called_once_with("sql")
+        # Default set hide_result=True
+        mock_after_exc.assert_called_once_with(hide_result=False)
+        expected_calls = [
+            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_cursor_execute")
     @mock.patch("osprofiler.sqlalchemy._after_cursor_execute")
     def test_disable_and_enable(self, mock_after_exc, mock_before_exc):