From 0f7ae6b2ba791908fa5d62d0de188a5bf89147a5 Mon Sep 17 00:00:00 2001
From: Itay Gibel <itay@heliosphere.io>
Date: Wed, 1 Sep 2021 10:40:14 +0300
Subject: [PATCH] Urllib3: extend request hook with request body and headers

---
 .../instrumentation/urllib3/__init__.py       | 30 +++++++++++++++++--
 .../tests/test_urllib3_integration.py         | 27 ++++++++++++++++-
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py b/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py
index 49d8a85b6b..078e295c34 100644
--- a/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py
+++ b/instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py
@@ -92,6 +92,18 @@ def response_hook(span, request, response):
 _RequestHookT = typing.Optional[
     typing.Callable[[Span, urllib3.connectionpool.HTTPConnectionPool], None]
 ]
+_ExtendedRequestHookT = typing.Optional[
+    typing.Callable[
+        [
+            Span,
+            urllib3.connectionpool.HTTPConnectionPool,
+            # Request headers dict
+            typing.Dict,
+            # Request Body
+            str,
+        ],
+        None]
+]
 _ResponseHookT = typing.Optional[
     typing.Callable[
         [
@@ -139,7 +151,7 @@ def _uninstrument(self, **kwargs):
 
 def _instrument(
     tracer,
-    request_hook: _RequestHookT = None,
+    request_hook: typing.Union[_RequestHookT, _ExtendedRequestHookT] = None,
     response_hook: _ResponseHookT = None,
     url_filter: _UrlFilterT = None,
 ):
@@ -150,6 +162,7 @@ def instrumented_urlopen(wrapped, instance, args, kwargs):
         method = _get_url_open_arg("method", args, kwargs).upper()
         url = _get_url(instance, args, kwargs, url_filter)
         headers = _prepare_headers(kwargs)
+        body = _get_url_open_arg("body", args, kwargs)
 
         span_name = "HTTP {}".format(method.strip())
         span_attributes = {
@@ -161,7 +174,7 @@ def instrumented_urlopen(wrapped, instance, args, kwargs):
             span_name, kind=SpanKind.CLIENT, attributes=span_attributes
         ) as span:
             if callable(request_hook):
-                request_hook(span, instance)
+                _call_request_hook(request_hook, span, instance, headers, body)
             inject(headers)
 
             with _suppress_further_instrumentation():
@@ -179,6 +192,19 @@ def instrumented_urlopen(wrapped, instance, args, kwargs):
     )
 
 
+def _call_request_hook(request_hook: typing.Union[_RequestHookT, _ExtendedRequestHookT],
+                       span: Span,
+                       connection_pool: urllib3.connectionpool.HTTPConnectionPool,
+                       headers: typing.Dict,
+                       body: str):
+    try:
+        # First assume request_hook is a function of type _ExtendedRequestHookT
+        request_hook(span, connection_pool, headers, body)
+    except TypeError:
+        # Fallback to call request_hook as a function of type _RequestHookT
+        request_hook(span, connection_pool)
+
+
 def _get_url_open_arg(name: str, args: typing.List, kwargs: typing.Mapping):
     arg_idx = _URL_OPEN_ARG_TO_INDEX_MAPPING.get(name)
     if arg_idx is not None:
diff --git a/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py b/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py
index 9e32162d34..c91420424f 100644
--- a/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py
+++ b/instrumentation/opentelemetry-instrumentation-urllib3/tests/test_urllib3_integration.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import json
 import typing
 from unittest import mock
 
@@ -279,3 +279,28 @@ def response_hook(span, request, response):
         self.assertEqual(span.name, "name set from hook")
         self.assertIn("response_hook_attr", span.attributes)
         self.assertEqual(span.attributes["response_hook_attr"], "value")
+
+    def test_extended_request_hook(self):
+        def extended_request_hook(span, request, headers, body):
+            span.set_attribute("request_hook_headers", json.dumps(headers))
+            span.set_attribute("request_hook_body", body)
+
+        URLLib3Instrumentor().uninstrument()
+        URLLib3Instrumentor().instrument(
+            request_hook=extended_request_hook,
+        )
+
+        headers = {"header1": "value1", "header2": "value2"}
+        body = "param1=1&param2=2"
+
+        pool = urllib3.HTTPConnectionPool("httpbin.org")
+        response = pool.request("GET", "/status/200", body=body, headers=headers)
+
+        self.assertEqual(b"Hello!", response.data)
+
+        span = self.assert_span()
+
+        self.assertIn("request_hook_headers", span.attributes)
+        self.assertEqual(span.attributes["request_hook_headers"], json.dumps(headers))
+        self.assertIn("request_hook_body", span.attributes)
+        self.assertEqual(span.attributes["request_hook_body"], body)