[requests] code refactoring in multiple modules

Author: Emanuele Palazzetti

ddtrace/contrib/requests/__init__.py

@@ -32,5 +32,11 @@
 
 with require_modules(required_modules) as missing_modules:
     if not missing_modules:
-        from .patch import TracedSession, patch, unpatch
-        __all__ = ['TracedSession', 'patch', 'unpatch']
+        from .patch import patch, unpatch
+        from .session import TracedSession
+
+        __all__ = [
+            'patch',
+            'unpatch',
+            'TracedSession',
+        ]

ddtrace/contrib/requests/connection.py

@@ -0,0 +1,56 @@
+import os
+import logging
+import ddtrace
+
+from ...ext import http
+from ...util import asbool
+from ...propagation.http import HTTPPropagator
+
+
+log = logging.getLogger(__name__)
+
+
+def _wrap_session_init(func, instance, args, kwargs):
+    """Configure tracing settings when the `Session` is initialized"""
+    func(*args, **kwargs)
+
+    # set tracer settings
+    distributed_tracing = asbool(os.environ.get('DATADOG_REQUESTS_DISTRIBUTED_TRACING')) or False
+    setattr(instance, 'distributed_tracing', distributed_tracing)
+
+
+def _wrap_request(func, instance, args, kwargs):
+    """Trace the `Session.request` instance method"""
+    tracer = getattr(instance, 'datadog_tracer', ddtrace.tracer)
+
+    # [TODO:christian] replace this with a unified way of handling options (eg, Pin)
+    distributed_tracing = getattr(instance, 'distributed_tracing', None)
+
+    # skip if tracing is not enabled
+    if not tracer.enabled:
+        return func(*args, **kwargs)
+
+    method = kwargs.get('method') or args[0]
+    url = kwargs.get('url') or args[1]
+    headers = kwargs.get('headers', {})
+
+    with tracer.trace("requests.request", span_type=http.TYPE) as span:
+        if distributed_tracing:
+            propagator = HTTPPropagator()
+            propagator.inject(span.context, headers)
+            kwargs['headers'] = headers
+
+        response = None
+        try:
+            response = func(*args, **kwargs)
+            return response
+        finally:
+            try:
+                span.set_tag(http.METHOD, method)
+                span.set_tag(http.URL, url)
+                if response is not None:
+                    span.set_tag(http.STATUS_CODE, response.status_code)
+                    # `span.error` must be an integer
+                    span.error = int(500 <= response.status_code)
+            except Exception:
+                log.debug("error patching tags", exc_info=True)

ddtrace/contrib/requests/patch.py

@@ -1,17 +1,9 @@
-import os
-import logging
-
-import wrapt
-import ddtrace
 import requests
 
-from ...ext import http
-from ...propagation.http import HTTPPropagator
-from ...utils.formats import asbool
-from ...utils.wrappers import unwrap as _u
-
+from wrapt import wrap_function_wrapper as _w
 
-log = logging.getLogger(__name__)
+from ...util import unwrap as _u
+from .connection import _wrap_session_init, _wrap_request
 
 
 def patch():
@@ -20,8 +12,8 @@ def patch():
         return
     setattr(requests, '__datadog_patch', True)
 
-    wrapt.wrap_function_wrapper('requests', 'Session.__init__', _session_initializer)
-    wrapt.wrap_function_wrapper('requests', 'Session.request', _traced_request_func)
+    _w('requests', 'Session.__init__', _wrap_session_init)
+    _w('requests', 'Session.request', _wrap_request)
 
 
 def unpatch():
@@ -32,70 +24,3 @@ def unpatch():
 
     _u(requests.Session, '__init__')
     _u(requests.Session, 'request')
-
-
-def _session_initializer(func, instance, args, kwargs):
-    """Define settings when requests client is initialized"""
-    func(*args, **kwargs)
-
-    # set tracer settings
-    distributed_tracing = asbool(os.environ.get('DATADOG_REQUESTS_DISTRIBUTED_TRACING')) or False
-    setattr(instance, 'distributed_tracing', distributed_tracing)
-
-
-def _traced_request_func(func, instance, args, kwargs):
-    """ traced_request is a tracing wrapper for requests' Session.request
-        instance method.
-    """
-
-    # perhaps a global tracer isn't what we want, so permit individual requests
-    # sessions to have their own (with the standard global fallback)
-    tracer = getattr(instance, 'datadog_tracer', ddtrace.tracer)
-
-    # [TODO:christian] replace this with a unified way of handling options (eg, Pin)
-    distributed_tracing = getattr(instance, 'distributed_tracing', None)
-
-    # bail on the tracing if not enabled.
-    if not tracer.enabled:
-        return func(*args, **kwargs)
-
-    method = kwargs.get('method') or args[0]
-    url = kwargs.get('url') or args[1]
-    headers = kwargs.get('headers', {})
-
-    with tracer.trace("requests.request", span_type=http.TYPE) as span:
-        if distributed_tracing:
-            propagator = HTTPPropagator()
-            propagator.inject(span.context, headers)
-            kwargs['headers'] = headers
-
-        resp = None
-        try:
-            resp = func(*args, **kwargs)
-            return resp
-        finally:
-            try:
-                _apply_tags(span, method, url, resp)
-            except Exception:
-                log.debug("error patching tags", exc_info=True)
-
-
-def _apply_tags(span, method, url, response):
-    """ apply_tags will patch the given span with tags about the given request. """
-    span.set_tag(http.METHOD, method)
-    span.set_tag(http.URL, url)
-    if response is not None:
-        span.set_tag(http.STATUS_CODE, response.status_code)
-        # `span.error` must be an integer
-        span.error = int(500 <= response.status_code)
-
-
-class TracedSession(requests.Session):
-    """ TracedSession is a requests' Session that is already patched.
-    """
-    pass
-
-
-# Always patch our traced session with the traced method (cheesy way of sharing
-# code)
-wrapt.wrap_function_wrapper(TracedSession, 'request', _traced_request_func)

ddtrace/contrib/requests/session.py

@@ -0,0 +1,18 @@
+import requests
+
+from wrapt import wrap_function_wrapper as _w
+
+from .connection import _wrap_session_init, _wrap_request
+
+
+class TracedSession(requests.Session):
+    """TracedSession is a requests' Session that is already traced.
+    You can use it if you want a finer grained control for your
+    HTTP clients.
+    """
+    pass
+
+
+# always patch our `TracedSession` when imported
+_w(TracedSession, 'request', _wrap_session_init)
+_w(TracedSession, 'request', _wrap_request)