Merge "Handle and report SQLAlchemy errors"1.14.0

2 files changed, 58 insertions, 7 deletions

diff --git a/osprofiler/sqlalchemy.py b/osprofiler/sqlalchemy.py
index 746b1b6..2118adc 100644
--- a/osprofiler/sqlalchemy.py
+++ b/osprofiler/sqlalchemy.py
@@ -14,9 +14,13 @@
 #    under the License.
 
 import contextlib
+import logging as log
+
+from oslo_utils import reflection
 
 from osprofiler import profiler
 
+LOG = log.getLogger(__name__)
 
 _DISABLED = False
 
@@ -44,6 +48,7 @@ def add_tracing(sqlalchemy, engine, name, hide_result=True):
             engine, "after_cursor_execute",
             _after_cursor_execute(hide_result=hide_result)
         )
+        sqlalchemy.event.listen(engine, "handle_error", handle_error)
 
 
 @contextlib.contextmanager
@@ -89,3 +94,22 @@ def _after_cursor_execute(hide_result=True):
             profiler.stop()
 
     return handler
+
+
+def handle_error(exception_context):
+    """Handle SQLAlchemy errors"""
+    exception_class_name = reflection.get_class_name(
+        exception_context.original_exception)
+    original_exception = str(exception_context.original_exception)
+    chained_exception = str(exception_context.chained_exception)
+
+    info = {
+        "etype": exception_class_name,
+        "db": {
+            "original_exception": original_exception,
+            "chained_exception": chained_exception
+        }
+    }
+    profiler.stop(info=info)
+    LOG.debug("OSProfiler has handled SQLAlchemy error: %s",
+              original_exception)
diff --git a/osprofiler/tests/unit/test_sqlalchemy.py b/osprofiler/tests/unit/test_sqlalchemy.py
index f7e2ead..0b20aa5 100644
--- a/osprofiler/tests/unit/test_sqlalchemy.py
+++ b/osprofiler/tests/unit/test_sqlalchemy.py
@@ -50,9 +50,30 @@ class SqlalchemyTracingTestCase(test.TestCase):
         }
         mock_profiler.stop.assert_called_once_with(info=info)
 
+    @mock.patch("osprofiler.sqlalchemy.profiler")
+    def test_handle_error(self, mock_profiler):
+        original_exception = Exception("error")
+        chained_exception = Exception("error and the reason")
+
+        sqlalchemy_exception_ctx = mock.MagicMock()
+        sqlalchemy_exception_ctx.original_exception = original_exception
+        sqlalchemy_exception_ctx.chained_exception = chained_exception
+
+        sqlalchemy.handle_error(sqlalchemy_exception_ctx)
+        expected_info = {
+            "etype": "Exception",
+            "db": {
+                "original_exception": str(original_exception),
+                "chained_exception": str(chained_exception),
+            }
+        }
+        mock_profiler.stop.assert_called_once_with(info=expected_info)
+
+    @mock.patch("osprofiler.sqlalchemy.handle_error")
     @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):
+    def test_add_tracing(self, mock_after_exc, mock_before_exc,
+                         mock_handle_error):
         sa = mock.MagicMock()
         engine = mock.MagicMock()
 
@@ -66,13 +87,16 @@ class SqlalchemyTracingTestCase(test.TestCase):
         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")
+            mock.call(engine, "after_cursor_execute", "after"),
+            mock.call(engine, "handle_error", mock_handle_error),
         ]
         self.assertEqual(sa.event.listen.call_args_list, expected_calls)
 
+    @mock.patch("osprofiler.sqlalchemy.handle_error")
     @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):
+    def test_wrap_session(self, mock_after_exc, mock_before_exc,
+                          mock_handle_error):
         sa = mock.MagicMock()
 
         @contextlib.contextmanager
@@ -96,16 +120,18 @@ class SqlalchemyTracingTestCase(test.TestCase):
         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")
+            mock.call(sess.bind, "after_cursor_execute", "after"),
+            mock.call(sess.bind, "handle_error", mock_handle_error),
         ]
 
         self.assertEqual(sa.event.listen.call_args_list, expected_calls)
 
+    @mock.patch("osprofiler.sqlalchemy.handle_error")
     @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):
+    def test_with_sql_result(self, mock_profiler, mock_after_exc,
+                             mock_before_exc, mock_handle_error):
         sa = mock.MagicMock()
         engine = mock.MagicMock()
 
@@ -119,7 +145,8 @@ class SqlalchemyTracingTestCase(test.TestCase):
         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")
+            mock.call(engine, "after_cursor_execute", "after"),
+            mock.call(engine, "handle_error", mock_handle_error),
         ]
         self.assertEqual(sa.event.listen.call_args_list, expected_calls)