Urllib3 request hook

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/open-telemetry/opentelemetry-python/compare/v1.5.0-0.24b0...HEAD)
 
+### Changed
+
+- `opentelemetry-instrumentation-urllib3` Added `_ExtendedRequestHookT` as an optional type of `request_hook`.
+   this type extends the existing `_RequestHookT` with two more fields - the request body and the request headers
+
 ## [1.5.0-0.24b0](https://github.com/open-telemetry/opentelemetry-python/releases/tag/v1.5.0-0.24b0) - 2021-08-26
 
 ### Added

instrumentation/opentelemetry-instrumentation-urllib3/src/opentelemetry/instrumentation/urllib3/__init__.py

@@ -92,6 +92,17 @@ 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,
+            typing.Dict,
+            typing.Optional[str],
+        ],
+        None,
+    ]
+]
 _ResponseHookT = typing.Optional[
     typing.Callable[
         [
@@ -139,7 +150,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 +161,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 +173,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 +191,21 @@ 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:

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
 
@@ -45,6 +45,7 @@ def setUp(self):
         httpretty.enable(allow_net_connect=False)
         httpretty.register_uri(httpretty.GET, self.HTTP_URL, body="Hello!")
         httpretty.register_uri(httpretty.GET, self.HTTPS_URL, body="Hello!")
+        httpretty.register_uri(httpretty.POST, self.HTTP_URL, body="Hello!")
 
     def tearDown(self):
         super().tearDown()
@@ -279,3 +280,30 @@ 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(
+            "POST", "/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)