Add post response hook

iopipe/agent.py

@@ -146,11 +146,11 @@ def wrapped(event, context):
             if timeout_duration > 0:
                 logger.debug("Setting timeout duration to %s" % timeout_duration)
 
-            result = None
+            response = None
 
             try:
                 with Timeout(timeout_duration):
-                    result = func(event, context)
+                    response = func(event, context)
             except Exception as e:
                 self.run_hooks("post:invoke", event=event, context=context)
 
@@ -169,6 +169,7 @@ def wrapped(event, context):
                 # event that a timeout occurs and an exception is subsequently raised
                 # within the handler
                 if self.report.sent is False:
+                    self.run_hooks("post:response", response=response)
                     self.report.prepare(e, frame)
                     self.run_hooks("pre:report")
                     self.report.send()
@@ -180,14 +181,15 @@ def wrapped(event, context):
 
                 if context.iopipe.disabled:
                     logger.debug("Reporting disabled for this invocation")
-                    return result
+                    return response
 
+                self.run_hooks("post:response", response=response)
                 self.report.prepare()
                 self.run_hooks("pre:report")
                 self.report.send()
                 self.run_hooks("post:report")
 
-                return result
+                return response
             finally:
                 self.wait_for_futures()
 
@@ -219,7 +221,7 @@ def instantiate(plugin):
 
         return loaded_plugins
 
-    def run_hooks(self, name, event=None, context=None):
+    def run_hooks(self, name, event=None, context=None, response=None):
         """
         Runs plugin hooks for each registered plugin.
         """
@@ -228,6 +230,7 @@ def run_hooks(self, name, event=None, context=None):
             "post:setup": lambda p: p.post_setup(self),
             "pre:invoke": lambda p: p.pre_invoke(event, context),
             "post:invoke": lambda p: p.post_invoke(event, context),
+            "post:response": lambda p: p.post_response(response),
             "pre:report": lambda p: p.pre_report(self.report),
             "post:report": lambda p: p.post_report(self.report),
         }

iopipe/contrib/eventinfo/plugin.py

@@ -41,6 +41,9 @@ def post_invoke(self, event, context):
         if self.enabled:
             metrics_for_event_type(event, context)
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         pass
 

iopipe/contrib/logger/plugin.py

@@ -110,6 +110,9 @@ def post_invoke(self, event, context):
         if self.enabled and self.redirect_stdout is True:
             sys.stdout = sys.__stdout__
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.enabled:
             self.handler.flush()

iopipe/contrib/profiler/plugin.py

@@ -63,6 +63,9 @@ def post_invoke(self, event, context):
             self.profile.disable()
             self.context.iopipe.label("@iopipe/plugin-profiler")
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.profile is not None:
             if self.signed_request is not None:

iopipe/contrib/trace/plugin.py

@@ -58,6 +58,9 @@ def post_invoke(self, event, context):
         if self.auto_http is True:
             restore_http_requests()
 
+    def post_response(self, response):
+        pass
+
     def pre_report(self, report):
         if self.auto_measure:
             add_timeline_measures(self.timeline)

iopipe/plugins.py

@@ -79,6 +79,10 @@ def pre_invoke(self, event, context):
     def post_invoke(self, event, context):
         return NotImplemented
 
+    @abc.abstractmethod
+    def post_response(self, response):
+        return NotImplemented
+
     @abc.abstractmethod
     def pre_report(self, report):
         return NotImplemented

tests/test_plugins.py

@@ -50,6 +50,9 @@ def pre_invoke(self, event, context):
         def post_invoke(self, event, context):
             pass
 
+        def post_response(self, response):
+            pass
+
         def pre_report(self, report):
             pass