Add an "inspect_error" middleware hook

Louis Opter · Louis Opter · commit e196f9b1cc6e · 2012-05-11T22:04:54.000-07:00
This middleware will be called when an exception is raised server side.

It receives the context of the call and the exception informations in
parameters. For example, this can be used to integrate Sentry and
ZeroRPC.
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
@@ -524,3 +524,43 @@ def echo(self, msg):
     assert subscriber_tracer._log == [
             ('load', publisher_tracer.trace_id),
             ]
+
+def test_inspect_error_middleware():
+
+    class InspectErrorMiddleware(Tracer):
+        def __init__(self):
+            self.called = False
+            Tracer.__init__(self, identity='[server]')
+
+        def inspect_error(self, task_context, exc_info):
+            assert 'trace_id' in task_context
+            exc_type, exc_value, exc_traceback = exc_info
+            self.called = True
+
+    class Srv(object):
+        def echo(self, msg):
+            raise RuntimeError(msg)
+
+    endpoint = random_ipc_endpoint()
+
+    middleware = InspectErrorMiddleware()
+    ctx = zerorpc.Context()
+    ctx.register_middleware(middleware)
+
+    module = Srv()
+    server = zerorpc.Server(module, context=ctx)
+    server.bind(endpoint)
+    gevent.spawn(server.run)
+
+    client = zerorpc.Client()
+    client.connect(endpoint)
+
+    try:
+        client.echo('This is a test which should call the InspectErrorMiddleware')
+    except zerorpc.exceptions.RemoteError as ex:
+        assert ex.name == 'RuntimeError'
+
+    client.close()
+    server.close()
+
+    assert middleware.called is True
diff --git a/zerorpc/context.py b/zerorpc/context.py
@@ -35,12 +35,13 @@ class Context(zmq.Context):
     def __init__(self):
         self._middlewares = []
         self._middlewares_hooks = {
-                'resolve_endpoint': [],
-                'raise_error': [],
-                'call_procedure': [],
-                'load_task_context': [],
-                'get_task_context': [],
-                }
+            'resolve_endpoint': [],
+            'raise_error': [],
+            'call_procedure': [],
+            'load_task_context': [],
+            'get_task_context': [],
+            'inspect_error': []
+        }
 
     @staticmethod
     def get_instance():
@@ -71,6 +72,12 @@ def middleware_resolve_endpoint(self, endpoint):
             endpoint = functor(endpoint)
         return endpoint
 
+    def middleware_inspect_error(self, exc_type, exc_value, exc_traceback):
+        exc_info = exc_type, exc_value, exc_traceback
+        task_context = self.middleware_get_task_context()
+        for functor in self._middlewares_hooks['inspect_error']:
+            functor(task_context, exc_info)
+
     def middleware_raise_error(self, event):
         for functor in self._middlewares_hooks['raise_error']:
             functor(event)
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -113,6 +113,9 @@ def _print_traceback(self, protocol_v1):
             traceback.print_exception(exc_type, exc_value, exc_traceback,
                     file=sys.stderr)
 
+            self._context.middleware_inspect_error(exc_type, exc_value,
+                    exc_traceback)
+
             if protocol_v1:
                 return (repr(exc_value),)
 
