[requests] migrate to the new Config system; remove __init__ wrap

Author: Emanuele Palazzetti

ddtrace/contrib/requests/connection.py

@@ -1,28 +1,17 @@
-import os
 import logging
 import ddtrace
 
+from ddtrace import config
+
 from .constants import DEFAULT_SERVICE
 
 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
-    service_name = os.environ.get('DATADOG_REQUESTS_SERVICE_NAME') or DEFAULT_SERVICE
-    setattr(instance, 'distributed_tracing', distributed_tracing)
-    setattr(instance, 'service_name', service_name)
-
-
 def _extract_service_name(session, span):
     """Extracts the right service name based on the following logic:
     - `requests` is the default service name
@@ -35,7 +24,7 @@ def _extract_service_name(session, span):
     The priority can be represented as:
     Updated service name > parent service name > default to `requests`.
     """
-    service_name = getattr(session, 'service_name', DEFAULT_SERVICE)
+    service_name = config.get_from(session)['service_name']
     if (service_name == DEFAULT_SERVICE and
             span._parent is not None and
             span._parent.service is not None):
@@ -45,11 +34,12 @@ def _extract_service_name(session, span):
 
 def _wrap_request(func, instance, args, kwargs):
     """Trace the `Session.request` instance method"""
+    # TODO[manu]: we already offer a way to provide the Global Tracer
+    # and is ddtrace.tracer; it's used only inside our tests and can
+    # be easily changed by providing a TracingTestCase that sets common
+    # tracing functionalities.
     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)
@@ -62,7 +52,8 @@ def _wrap_request(func, instance, args, kwargs):
         # update the span service name before doing any action
         span.service = _extract_service_name(instance, span)
 
-        if distributed_tracing:
+        # propagate distributed tracing headers
+        if config.get_from(instance)['distributed_tracing']:
             propagator = HTTPPropagator()
             propagator.inject(span.context, headers)
             kwargs['headers'] = headers

ddtrace/contrib/requests/legacy.py

@@ -0,0 +1,11 @@
+from ddtrace import config
+
+
+def _distributed_tracing(self):
+    """Backward compatibility"""
+    return config.get_from(self)['distributed_tracing']
+
+
+def _distributed_tracing_setter(self, value):
+    """Backward compatibility"""
+    config.get_from(self)['distributed_tracing'] = value

ddtrace/contrib/requests/patch.py

@@ -2,8 +2,20 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from ...util import unwrap as _u
-from .connection import _wrap_session_init, _wrap_request
+from ddtrace import config
+from ddtrace.pin import Pin
+
+from ...util import asbool, get_env, unwrap as _u
+from .legacy import _distributed_tracing, _distributed_tracing_setter
+from .constants import DEFAULT_SERVICE
+from .connection import _wrap_request
+
+
+# requests default settings
+config._add('requests',{
+    'service_name': get_env('requests', 'service_name', DEFAULT_SERVICE),
+    'distributed_tracing': asbool(get_env('requests', 'distributed_tracing', False)),
+})
 
 
 def patch():
@@ -12,8 +24,18 @@ def patch():
         return
     setattr(requests, '__datadog_patch', True)
 
-    _w('requests', 'Session.__init__', _wrap_session_init)
     _w('requests', 'Session.request', _wrap_request)
+    Pin(
+        service=config.requests['service_name'],
+        _config=config.requests,
+    ).onto(requests.Session)
+
+    # [Backward compatibility]: `session.distributed_tracing` should point and
+    # update the `Pin` configuration instead. This block adds a property so that
+    # old implementations work as expected
+    fn = property(_distributed_tracing)
+    fn = fn.setter(_distributed_tracing_setter)
+    requests.Session.distributed_tracing = fn
 
 
 def unpatch():
@@ -22,5 +44,4 @@ def unpatch():
         return
     setattr(requests, '__datadog_patch', False)
 
-    _u(requests.Session, '__init__')
     _u(requests.Session, 'request')

ddtrace/contrib/requests/session.py

@@ -2,7 +2,7 @@
 
 from wrapt import wrap_function_wrapper as _w
 
-from .connection import _wrap_session_init, _wrap_request
+from .connection import _wrap_request
 
 
 class TracedSession(requests.Session):
@@ -14,5 +14,4 @@ class TracedSession(requests.Session):
 
 
 # always patch our `TracedSession` when imported
-_w(TracedSession, 'request', _wrap_session_init)
 _w(TracedSession, 'request', _wrap_request)

tests/contrib/requests/test_requests.py

@@ -1,8 +1,9 @@
 import unittest
 
 from requests import Session
-from nose.tools import eq_, assert_raises
+from nose.tools import eq_
 
+from ddtrace import config
 from ddtrace.ext import http, errors
 from ddtrace.contrib.requests import patch, unpatch
 
@@ -153,7 +154,8 @@ def test_default_service_name(self):
 
     def test_user_set_service_name(self):
         # ensure a service name set by the user has precedence
-        self.session.service_name = 'clients'
+        cfg = config.get_from(self.session)
+        cfg['service_name'] = 'clients'
         out = self.session.get(URL_200)
         eq_(out.status_code, 200)
 
@@ -194,8 +196,9 @@ def test_parent_without_service_name(self):
     def test_user_service_name_precedence(self):
         # ensure the user service name takes precedence over
         # the parent Span
+        cfg = config.get_from(self.session)
+        cfg['service_name'] = 'clients'
         with self.tracer.trace('parent.span', service='web'):
-            self.session.service_name = 'clients'
             out = self.session.get(URL_200)
             eq_(out.status_code, 200)
 

tests/contrib/requests/test_requests_distributed.py

@@ -1,6 +1,8 @@
 from requests_mock import Adapter
 from nose.tools import eq_, assert_in, assert_not_in
 
+from ddtrace import config
+
 from .test_requests import BaseRequestTestCase
 
 
@@ -24,10 +26,26 @@ def headers_not_here(self, tracer, request):
         assert_not_in('x-datadog-parent-id', headers)
         return True
 
+    def test_propagation_default(self):
+        # ensure by default, distributed tracing is disabled
+        adapter = Adapter()
+        self.session.mount('mock', adapter)
+
+        with self.tracer.trace('root'):
+            def matcher(request):
+                return self.headers_not_here(self.tracer, request)
+            adapter.register_uri('GET', 'mock://datadog/foo', additional_matcher=matcher, text='bar')
+            resp = self.session.get('mock://datadog/foo')
+            eq_(200, resp.status_code)
+            eq_('bar', resp.text)
+
     def test_propagation_true(self):
+        # [Backward compatibility]: ensure users can switch the distributed
+        # tracing flag using the `Session` attribute
+        cfg = config.get_from(self.session)
+        cfg['distributed_tracing'] = True
         adapter = Adapter()
         self.session.mount('mock', adapter)
-        self.session.distributed_tracing = True
 
         with self.tracer.trace('root') as root:
             def matcher(request):
@@ -45,6 +63,46 @@ def matcher(request):
         eq_(root.span_id, req.parent_id)
 
     def test_propagation_false(self):
+        # [Backward compatibility]: ensure users can switch the distributed
+        # tracing flag using the `Session` attribute
+        cfg = config.get_from(self.session)
+        cfg['distributed_tracing'] = False
+        adapter = Adapter()
+        self.session.mount('mock', adapter)
+
+        with self.tracer.trace('root'):
+            def matcher(request):
+                return self.headers_not_here(self.tracer, request)
+            adapter.register_uri('GET', 'mock://datadog/foo', additional_matcher=matcher, text='bar')
+            resp = self.session.get('mock://datadog/foo')
+            eq_(200, resp.status_code)
+            eq_('bar', resp.text)
+
+    def test_propagation_true_legacy(self):
+        # [Backward compatibility]: ensure users can switch the distributed
+        # tracing flag using the `Session` attribute
+        adapter = Adapter()
+        self.session.mount('mock', adapter)
+        self.session.distributed_tracing = True
+
+        with self.tracer.trace('root') as root:
+            def matcher(request):
+                return self.headers_here(self.tracer, request, root)
+            adapter.register_uri('GET', 'mock://datadog/foo', additional_matcher=matcher, text='bar')
+            resp = self.session.get('mock://datadog/foo')
+            eq_(200, resp.status_code)
+            eq_('bar', resp.text)
+
+        spans = self.tracer.writer.spans
+        root, req = spans
+        eq_('root', root.name)
+        eq_('requests.request', req.name)
+        eq_(root.trace_id, req.trace_id)
+        eq_(root.span_id, req.parent_id)
+
+    def test_propagation_false_legacy(self):
+        # [Backward compatibility]: ensure users can switch the distributed
+        # tracing flag using the `Session` attribute
         adapter = Adapter()
         self.session.mount('mock', adapter)
         self.session.distributed_tracing = False