Limit HTTP headers collected, allow overide (#326)

* Limit HTTP headers collected, allow overide

Closes #274 #305

* Add docs to README

README.md

@@ -456,6 +456,16 @@ def http_filter(request, response):
 iopipe = IOpipe(plugins=[TracePlugin(auto_http=True, http_filter=http_filter)])
 ```
 
+To add additional HTTP headers to your ttrace data use `http_headers`:
+
+```python
+http_headers = ['Cache-Control', 'Etag']
+
+iopipe = IOpipe(plugins=[TracePlugin(auto_http=True, http_headers=http_headers)
+```
+
+## Plugins
+
 ### Creating Plugins
 
 To create an IOpipe plugin you must implement the `iopipe.plugins.Plugin` interface.

iopipe/contrib/trace/auto_http.py

@@ -16,7 +16,7 @@
 except ImportError:
     BotocoreVendoredSession = None
 
-from iopipe.compat import urlparse
+from iopipe.compat import string_types, urlparse
 from .util import ensure_utf8
 
 if RequestsSession is not None:
@@ -29,21 +29,12 @@
     original_botocore_vendored_session_send = BotocoreVendoredSession.send
 
 INCLUDE_HEADERS = [
-    "accept",
-    "accept-encoding",
-    "age",
-    "cache-control",
-    "connection",
-    "content-encoding",
     "content-length",
     "content-type",
-    "date",
-    "etag",
     "host",
     "server",
-    "strict-transport-security",
     "user-agent",
-    "vary",
+    "x-amz-target",
 ]
 
 Request = collections.namedtuple(
@@ -67,7 +58,7 @@
 )
 
 
-def patch_requests_session_send(context, http_filter):
+def patch_requests_session_send(context, http_filter, http_headers):
     """
     Monkey patches requests' session class, if available. Overloads the
     send method to add tracing and metrics collection.
@@ -85,15 +76,15 @@ def send(self, *args, **kwargs):
         trace = context.iopipe.mark.measure(id)
         context.iopipe.mark.delete(id)
         collect_metrics_for_response(
-            response.request, response, context, trace, http_filter
+            response.request, response, context, trace, http_filter, http_headers
         )
         return response
 
     RequestsSession.send = send
     RequestsSession.__monkey_patched = True
 
 
-def patch_botocore_session_send(context, http_filter):
+def patch_botocore_session_send(context, http_filter, http_headers):
     """
     Monkey patches botocore's session, if available. Overloads the
     session class' send method to add tracing and metric collection.
@@ -110,14 +101,16 @@ def send(self, *args, **kwargs):
             response = original_botocore_session_send(self, *args, **kwargs)
         trace = context.iopipe.mark.measure(id)
         context.iopipe.mark.delete(id)
-        collect_metrics_for_response(args[0], response, context, trace, http_filter)
+        collect_metrics_for_response(
+            args[0], response, context, trace, http_filter, http_headers
+        )
         return response
 
     BotocoreSession.send = send
     BotocoreSession.__monkey_patched = True
 
 
-def patch_botocore_vendored_session_send(context, http_filter):
+def patch_botocore_vendored_session_send(context, http_filter, http_headers):
     """
     Monkey patches botocore's vendored requests, if available. Overloads the
     session class' send method to add tracing and metric collection.
@@ -135,7 +128,7 @@ def send(self, *args, **kwargs):
         trace = context.iopipe.mark.measure(id)
         context.iopipe.mark.delete(id)
         collect_metrics_for_response(
-            response.request, response, context, trace, http_filter
+            response.request, response, context, trace, http_filter, http_headers
         )
         return response
 
@@ -164,25 +157,33 @@ def restore_botocore_vendored_session_send():
         delattr(BotocoreVendoredSession, "__monkey_patched")
 
 
-def patch_requests(context, http_filter):
-    patch_requests_session_send(context, http_filter)
-    patch_botocore_session_send(context, http_filter)
-    patch_botocore_vendored_session_send(context, http_filter)
+def patch_http_requests(context, http_filter, http_headers):
+    patch_requests_session_send(context, http_filter, http_headers)
+    patch_botocore_session_send(context, http_filter, http_headers)
+    patch_botocore_vendored_session_send(context, http_filter, http_headers)
 
 
-def restore_requests():
+def restore_http_requests():
     restore_requests_session_send()
     restore_botocore_session_send()
     restore_botocore_vendored_session_send()
 
 
 def collect_metrics_for_response(
-    http_request, http_response, context, trace, http_filter
+    http_request, http_response, context, trace, http_filter, http_headers
 ):
     """
     Collects relevant metrics from a requests Response object and adds them to
     the IOpipe context.
     """
+    include_headers = INCLUDE_HEADERS
+    if isinstance(http_headers, (list, tuple)):
+        include_headers = include_headers + [
+            key.lower()
+            for key in http_headers
+            if isinstance(http_headers, string_types)
+        ]
+
     request = None
     if http_request:
         parsed_url = None
@@ -194,7 +195,7 @@ def collect_metrics_for_response(
             request_headers = [
                 {"key": ensure_utf8(k), "string": ensure_utf8(v)}
                 for k, v in http_request.headers.items()
-                if k.lower() in INCLUDE_HEADERS
+                if k.lower() in include_headers
             ]
 
         request = Request(
@@ -219,7 +220,7 @@ def collect_metrics_for_response(
         response_headers = [
             {"key": ensure_utf8(k), "string": ensure_utf8(v)}
             for k, v in http_response.headers.items()
-            if k.lower() in INCLUDE_HEADERS
+            if k.lower() in include_headers
         ]
 
     response = Response(

iopipe/contrib/trace/plugin.py

@@ -3,7 +3,7 @@
 
 from iopipe.plugins import Plugin
 
-from .auto_http import patch_requests, restore_requests
+from .auto_http import patch_http_requests, restore_http_requests
 from .marker import Marker
 from .timeline import Timeline
 from .util import add_timeline_measures
@@ -15,7 +15,9 @@ class TracePlugin(Plugin):
     homepage = "https://github.com/iopipe/iopipe-python#trace-plugin"
     enabled = True
 
-    def __init__(self, auto_measure=True, auto_http=True, http_filter=None):
+    def __init__(
+        self, auto_measure=True, auto_http=True, http_filter=None, http_headers=None
+    ):
         """
         Instantiates the trace plugin
 
@@ -25,12 +27,15 @@ def __init__(self, auto_measure=True, auto_http=True, http_filter=None):
         :type auto_http: bool
         :param http_filter: A callable to filter http requests
         :type http_filter: function
+        :param http_headers: Additional HTTP headers to collect
+        :type http_headers: list|tuple
         """
         self.auto_measure = auto_measure
         self.auto_http = auto_http is True or strtobool(
             os.getenv("IOPIPE_TRACE_AUTO_HTTP_ENABLED", "false")
         )
         self.http_filter = http_filter
+        self.http_headers = http_headers
 
         self.timeline = Timeline()
 
@@ -45,13 +50,13 @@ def pre_invoke(self, event, context):
         context.iopipe.register("mark", Marker(self.timeline, context))
 
         if self.auto_http is True:
-            patch_requests(context, self.http_filter)
+            patch_http_requests(context, self.http_filter, self.http_headers)
 
     def post_invoke(self, event, context):
         context.iopipe.unregister("mark")
 
         if self.auto_http is True:
-            restore_requests()
+            restore_http_requests()
 
     def pre_report(self, report):
         if self.auto_measure:

tests/contrib/trace/test_auto_http.py

@@ -2,7 +2,7 @@
 from botocore.vendored.requests.sessions import Session as BotocoreVendoredSession
 from requests.sessions import Session as RequestsSession
 
-from iopipe.contrib.trace.auto_http import patch_requests, restore_requests
+from iopipe.contrib.trace.auto_http import patch_http_requests, restore_http_requests
 
 
 def test_monkey_patching(mock_context):
@@ -13,13 +13,15 @@ def test_monkey_patching(mock_context):
     def mock_filter(request, response):
         return request, response
 
-    patch_requests(mock_context, mock_filter)
+    http_headers = ["Cache-Control"]
+
+    patch_http_requests(mock_context, mock_filter, http_headers)
 
     assert hasattr(RequestsSession, "__monkey_patched")
     assert hasattr(BotocoreSession, "__monkey_patched")
     assert hasattr(BotocoreVendoredSession, "__monkey_patched")
 
-    restore_requests()
+    restore_http_requests()
 
     assert not hasattr(RequestsSession, "__monkey_patched")
     assert not hasattr(BotocoreSession, "__monkey_patched")