[requests] add unpatch and double-patch protection (#404)

Author: Emanuele Palazzetti

ddtrace/contrib/requests/__init__.py

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

ddtrace/contrib/requests/patch.py

@@ -1,30 +1,43 @@
-"""
-Tracing for the requests library.
-
-https://github.com/kennethreitz/requests
-"""
-
-# stdlib
 import logging
 
-# 3p
-import requests
 import wrapt
+import requests
 
-# project
 import ddtrace
-from ddtrace.ext import http
+
+from ...ext import http
 from ...propagation.http import HTTPPropagator
+from ...util import unwrap as _u
 
 
 log = logging.getLogger(__name__)
 
 
 def patch():
-    """ Monkeypatch the requests library to trace http calls. """
+    """Activate http calls tracing"""
+    if getattr(requests, '__datadog_patch', False):
+        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)
 
 
+def unpatch():
+    """Disable traced sessions"""
+    if not getattr(requests, '__datadog_patch', False):
+        return
+    setattr(requests, '__datadog_patch', False)
+
+    _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)
+
+
 def _traced_request_func(func, instance, args, kwargs):
     """ traced_request is a tracing wrapper for requests' Session.request
         instance method.

tests/contrib/requests/test_requests.py

@@ -1,105 +1,125 @@
+import unittest
 
-# 3p
-from nose.tools import eq_, assert_raises
 from requests import Session
+from nose.tools import eq_, assert_raises
 
-# project
-from ddtrace.contrib.requests import TracedSession
 from ddtrace.ext import http, errors
-from tests.test_tracer import get_dummy_tracer
+from ddtrace.contrib.requests import patch, unpatch
+
+from ...test_tracer import get_dummy_tracer
 
 # socket name comes from https://english.stackexchange.com/a/44048
 SOCKET = 'httpbin.org'
 URL_200 = 'http://{}/status/200'.format(SOCKET)
 URL_500 = 'http://{}/status/500'.format(SOCKET)
 
-class TestRequests(object):
 
-    @staticmethod
-    def test_resource_path():
-        tracer, session = get_traced_session()
-        out = session.get(URL_200)
+class BaseRequestTestCase(unittest.TestCase):
+    """Create a traced Session, patching during the setUp and
+    unpatching after the tearDown
+    """
+    def setUp(self):
+        patch()
+        self.tracer = get_dummy_tracer()
+        self.session = Session()
+        setattr(self.session, 'datadog_tracer', self.tracer)
+
+    def tearDown(self):
+        unpatch()
+
+
+class TestRequests(BaseRequestTestCase):
+    def test_resource_path(self):
+        out = self.session.get(URL_200)
         eq_(out.status_code, 200)
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 1)
         s = spans[0]
         eq_(s.get_tag("http.url"), URL_200)
 
-    @staticmethod
-    def test_tracer_disabled():
+    def test_tracer_disabled(self):
         # ensure all valid combinations of args / kwargs work
-        tracer, session = get_traced_session()
-        tracer.enabled = False
-        out = session.get(URL_200)
+        self.tracer.enabled = False
+        out = self.session.get(URL_200)
         eq_(out.status_code, 200)
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 0)
 
-    @staticmethod
-    def test_args_kwargs():
+    def test_args_kwargs(self):
         # ensure all valid combinations of args / kwargs work
-        tracer, session = get_traced_session()
         url = URL_200
         method = 'GET'
         inputs = [
                 ([], {'method': method, 'url': url}),
                 ([method], {'url': url}),
                 ([method, url], {}),
         ]
-        untraced = Session()
+
         for args, kwargs in inputs:
-            # ensure an untraced request works with these args
-            out = untraced.request(*args, **kwargs)
-            eq_(out.status_code, 200)
-            out = session.request(*args, **kwargs)
+            # ensure a traced request works with these args
+            out = self.session.request(*args, **kwargs)
             eq_(out.status_code, 200)
             # validation
-            spans = tracer.writer.pop()
+            spans = self.tracer.writer.pop()
             eq_(len(spans), 1)
             s = spans[0]
             eq_(s.get_tag(http.METHOD), 'GET')
             eq_(s.get_tag(http.STATUS_CODE), '200')
 
+    def test_untraced_request(self):
+        # ensure the unpatch removes tracing
+        unpatch()
+        untraced = Session()
+
+        out = untraced.get(URL_200)
+        eq_(out.status_code, 200)
+        # validation
+        spans = self.tracer.writer.pop()
+        eq_(len(spans), 0)
+
+    def test_double_patch(self):
+        # ensure that double patch doesn't duplicate instrumentation
+        patch()
+        session = Session()
+        setattr(session, 'datadog_tracer', self.tracer)
 
-    @staticmethod
-    def test_200():
-        tracer, session = get_traced_session()
         out = session.get(URL_200)
         eq_(out.status_code, 200)
+        spans = self.tracer.writer.pop()
+        eq_(len(spans), 1)
+
+    def test_200(self):
+        out = self.session.get(URL_200)
+        eq_(out.status_code, 200)
         # validation
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 1)
         s = spans[0]
         eq_(s.get_tag(http.METHOD), 'GET')
         eq_(s.get_tag(http.STATUS_CODE), '200')
         eq_(s.error, 0)
         eq_(s.span_type, http.TYPE)
 
-    @staticmethod
-    def test_post_500():
-        tracer, session = get_traced_session()
-        out = session.post(URL_500)
+    def test_post_500(self):
+        out = self.session.post(URL_500)
         # validation
         eq_(out.status_code, 500)
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 1)
         s = spans[0]
         eq_(s.get_tag(http.METHOD), 'POST')
         eq_(s.get_tag(http.STATUS_CODE), '500')
         eq_(s.error, 1)
 
-    @staticmethod
-    def test_non_existant_url():
-        tracer, session = get_traced_session()
-
+    def test_non_existant_url(self):
         try:
-            session.get('http://doesnotexist.google.com')
+            self.session.get('http://doesnotexist.google.com')
         except Exception:
             pass
         else:
             assert 0, "expected error"
 
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 1)
         s = spans[0]
         eq_(s.get_tag(http.METHOD), 'GET')
@@ -109,23 +129,13 @@ def test_non_existant_url():
         assert "Traceback (most recent call last)" in s.get_tag(errors.STACK)
         assert "requests.exception" in s.get_tag(errors.TYPE)
 
-
-    @staticmethod
-    def test_500():
-        tracer, session = get_traced_session()
-        out = session.get(URL_500)
+    def test_500(self):
+        out = self.session.get(URL_500)
         eq_(out.status_code, 500)
 
-        spans = tracer.writer.pop()
+        spans = self.tracer.writer.pop()
         eq_(len(spans), 1)
         s = spans[0]
         eq_(s.get_tag(http.METHOD), 'GET')
         eq_(s.get_tag(http.STATUS_CODE), '500')
         eq_(s.error, 1)
-
-
-def get_traced_session():
-    tracer = get_dummy_tracer()
-    session = TracedSession()
-    setattr(session, 'datadog_tracer', tracer)
-    return tracer, session

tests/contrib/requests/test_requests_distributed.py

@@ -1,13 +1,10 @@
-
-# 3p
-from nose.tools import eq_, assert_in, assert_not_in
 from requests_mock import Adapter
+from nose.tools import eq_, assert_in, assert_not_in
 
-# project
-from .test_requests import get_traced_session
+from .test_requests import BaseRequestTestCase
 
-class TestRequestsDistributed(object):
 
+class TestRequestsDistributed(BaseRequestTestCase):
     def headers_here(self, tracer, request, root_span):
         # Use an additional matcher to query the request headers.
         # This is because the parent_id can only been known within such a callback,
@@ -29,19 +26,18 @@ def headers_not_here(self, tracer, request):
 
     def test_propagation_true(self):
         adapter = Adapter()
-        tracer, session = get_traced_session()
-        session.mount('mock', adapter)
-        session.distributed_tracing = True
+        self.session.mount('mock', adapter)
+        self.session.distributed_tracing = True
 
-        with tracer.trace('root') as root:
+        with self.tracer.trace('root') as root:
             def matcher(request):
-                return self.headers_here(tracer, request, root)
+                return self.headers_here(self.tracer, request, root)
             adapter.register_uri('GET', 'mock://datadog/foo', additional_matcher=matcher, text='bar')
-            resp = session.get('mock://datadog/foo')
+            resp = self.session.get('mock://datadog/foo')
             eq_(200, resp.status_code)
             eq_('bar', resp.text)
 
-        spans = tracer.writer.spans
+        spans = self.tracer.writer.spans
         root, req = spans
         eq_('root', root.name)
         eq_('requests.request', req.name)
@@ -50,14 +46,13 @@ def matcher(request):
 
     def test_propagation_false(self):
         adapter = Adapter()
-        tracer, session = get_traced_session()
-        session.mount('mock', adapter)
-        session.distributed_tracing = False
+        self.session.mount('mock', adapter)
+        self.session.distributed_tracing = False
 
-        with tracer.trace('root'):
+        with self.tracer.trace('root'):
             def matcher(request):
-                return self.headers_not_here(tracer, request)
+                return self.headers_not_here(self.tracer, request)
             adapter.register_uri('GET', 'mock://datadog/foo', additional_matcher=matcher, text='bar')
-            resp = session.get('mock://datadog/foo')
+            resp = self.session.get('mock://datadog/foo')
             eq_(200, resp.status_code)
             eq_('bar', resp.text)