Add option for specifying extra headers in RequestsWrapper (#4208)

datadog_checks_base/datadog_checks/base/stubs/datadog_agent.py

@@ -12,7 +12,11 @@ def log(*args, **kwargs):
 
 
 def get_config(*args, **kwargs):
-    return ""
+    return ''
+
+
+def get_version():
+    return '0.0.0'
 
 
 def warning(msg, *args, **kwargs):

datadog_checks_base/datadog_checks/base/utils/headers.py

@@ -1,10 +1,30 @@
 # (C) Datadog, Inc. 2018
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
+from collections import OrderedDict
+
+from six import iteritems
+
 try:
     import datadog_agent
 except ImportError:
-    datadog_agent = None
+    from ..stubs import datadog_agent
+
+
+def get_default_headers():
+    # http://docs.python-requests.org/en/master/user/advanced/#header-ordering
+    # TODO: Use built-in when we drop Python 2 as dictionaries are guaranteed to be ordered in Python 3.6+ (and PyPy)
+    return OrderedDict(
+        (
+            # Default to `0.0.0` if no version is found
+            ('User-Agent', 'Datadog Agent/{}'.format(datadog_agent.get_version() or '0.0.0')),
+        )
+    )
+
+
+def update_headers(headers, extra_headers):
+    # Ensure the values are strings
+    headers.update((key, str(value)) for key, value in iteritems(extra_headers))
 
 
 def headers(agentConfig, **kwargs):

datadog_checks_base/datadog_checks/base/utils/http.py

@@ -5,7 +5,6 @@
 import os
 import threading
 import warnings
-from collections import OrderedDict
 from contextlib import contextmanager
 
 import requests
@@ -15,6 +14,7 @@
 
 from ..config import is_affirmative
 from ..errors import ConfigurationError
+from .headers import get_default_headers, update_headers
 
 try:
     from contextlib import ExitStack
@@ -32,10 +32,14 @@
 
 LOGGER = logging.getLogger(__file__)
 
+# The timeout should be slightly larger than a multiple of 3,
+# which is the default TCP packet retransmission window. See:
+# https://tools.ietf.org/html/rfc2988
 DEFAULT_TIMEOUT = 10
 
 STANDARD_FIELDS = {
     'connect_timeout': None,
+    'extra_headers': None,
     'headers': None,
     'kerberos_auth': None,
     'kerberos_delegate': False,
@@ -147,9 +151,13 @@ def __init__(self, instance, init_config, remapper=None, logger=None):
 
         # http://docs.python-requests.org/en/master/user/quickstart/#custom-headers
         # http://docs.python-requests.org/en/master/user/advanced/#header-ordering
-        headers = None
+        headers = get_default_headers()
         if config['headers']:
-            headers = OrderedDict((key, str(value)) for key, value in iteritems(config['headers']))
+            headers.clear()
+            update_headers(headers, config['headers'])
+
+        if config['extra_headers']:
+            update_headers(headers, config['extra_headers'])
 
         # http://docs.python-requests.org/en/master/user/authentication/
         auth = None

datadog_checks_base/tests/test_http.py

@@ -66,7 +66,7 @@ def test_config_default(self):
         init_config = {}
         http = RequestsWrapper(instance, init_config)
 
-        assert http.options['headers'] is None
+        assert http.options['headers'] == {'User-Agent': 'Datadog Agent/0.0.0'}
 
     def test_config_headers(self):
         headers = OrderedDict((('key1', 'value1'), ('key2', 'value2')))
@@ -83,6 +83,23 @@ def test_config_headers_string_values(self):
 
         assert http.options['headers'] == {'answer': '42'}
 
+    def test_config_extra_headers(self):
+        headers = OrderedDict((('key1', 'value1'), ('key2', 'value2')))
+        instance = {'extra_headers': headers}
+        init_config = {}
+        http = RequestsWrapper(instance, init_config)
+
+        complete_headers = OrderedDict({'User-Agent': 'Datadog Agent/0.0.0'})
+        complete_headers.update(headers)
+        assert list(iteritems(http.options['headers'])) == list(iteritems(complete_headers))
+
+    def test_config_extra_headers_string_values(self):
+        instance = {'extra_headers': {'answer': 42}}
+        init_config = {}
+        http = RequestsWrapper(instance, init_config)
+
+        assert http.options['headers'] == {'User-Agent': 'Datadog Agent/0.0.0', 'answer': '42'}
+
 
 class TestVerify:
     def test_config_default(self):