From aaaf8930ce405bc376d7293027684adf82c5e60c Mon Sep 17 00:00:00 2001
From: Joshua Wise <joshua@joshuawise.com>
Date: Mon, 2 Mar 2020 23:32:51 -0500
Subject: [PATCH] patch: add urllib patch

---
 beeline/patch/urllib.py | 51 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 beeline/patch/urllib.py

diff --git a/beeline/patch/urllib.py b/beeline/patch/urllib.py
new file mode 100644
index 0000000..e11b160
--- /dev/null
+++ b/beeline/patch/urllib.py
@@ -0,0 +1,51 @@
+from wrapt import wrap_function_wrapper
+import beeline
+import urllib.request
+
+def _urllibopen(_urlopen, instance, args, kwargs):
+    if type(args[0]) != urllib.request.Request:
+        args[0] = urllib.request.Request(args[0])
+    
+    span = beeline.start_span(context={"meta.type": "http_client"})
+    
+    b = beeline.get_beeline()
+    if b:
+        context = b.tracer_impl.marshal_trace_context()
+        if context:
+            b.log("urllib lib - adding trace context to outbound request: %s", context)
+            args[0].headers['X-Honeycomb-Trace'] = context
+        else:
+            b.log("urllib lib - no trace context found")
+    
+    try:
+        resp = None
+        beeline.add_context({
+            "name": "urllib_%s" % args[0].get_method(),
+            "request.method": args[0].get_method(),
+            "request.uri": args[0].full_url
+        })
+        resp = _urlopen(*args, **kwargs)
+        return resp
+    except Exception as e:
+        beeline.add_context({
+            "request.error_type": str(type(e)),
+            "request.error": beeline.internal.stringify_exception(e),
+        })
+        raise
+    finally:
+        if resp:
+            beeline.add_context_field("response.status_code", resp.status)
+            content_type = resp.getheader('content-type')
+            if content_type:
+                beeline.add_context_field("response.content_type", content_type)
+            content_length = resp.getheader('content-length')
+            if content_length:
+                beeline.add_context_field("response.content_length", content_length)
+            
+        beeline.finish_span(span)
+
+# Note that this only patches urllib.request.urlopen, not
+# http.client.HTTPConnection.  The latter is a lot more of a pain to figure
+# out what, exactly, the lifetime of the span ought to be -- but most people
+# who plan to block and do nothing else use urlopen, anyway.
+wrap_function_wrapper('urllib.request', 'urlopen', _urllibopen)