From a3cc53e2d45616ba067bbfdc9d7510f08456ffa1 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Fri, 7 Jun 2019 18:40:01 -0700 Subject: [PATCH 1/2] Limit HTTP headers collected, allow overide Closes #274 #305 --- iopipe/contrib/trace/auto_http.py | 51 ++++++++++++++------------- iopipe/contrib/trace/plugin.py | 13 ++++--- tests/contrib/trace/test_auto_http.py | 8 +++-- 3 files changed, 40 insertions(+), 32 deletions(-) diff --git a/iopipe/contrib/trace/auto_http.py b/iopipe/contrib/trace/auto_http.py index 658c949e..f7a98aa8 100644 --- a/iopipe/contrib/trace/auto_http.py +++ b/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,7 +76,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 @@ -93,7 +84,7 @@ def send(self, *args, **kwargs): 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( diff --git a/iopipe/contrib/trace/plugin.py b/iopipe/contrib/trace/plugin.py index 2bf46b7c..bfb64e0f 100644 --- a/iopipe/contrib/trace/plugin.py +++ b/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: diff --git a/tests/contrib/trace/test_auto_http.py b/tests/contrib/trace/test_auto_http.py index 5f86d2fd..ead39eea 100644 --- a/tests/contrib/trace/test_auto_http.py +++ b/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") From 7d9445cf419b1748b94d497625c20ae3ed108109 Mon Sep 17 00:00:00 2001 From: Michael Lavers Date: Mon, 10 Jun 2019 09:59:05 -0700 Subject: [PATCH 2/2] Add docs to README --- README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/README.md b/README.md index ee684e0e..43f126d0 100644 --- a/README.md +++ b/README.md @@ -423,6 +423,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.