Ensure a custom User-Agent header is not overwritten (#992)

elasticsearch/connection/base.py

@@ -1,11 +1,14 @@
 import logging
 
+from platform import python_version
+
 try:
     import simplejson as json
 except ImportError:
     import json
 
 from ..exceptions import TransportError, HTTP_EXCEPTIONS
+from .. import __versionstr__
 
 logger = logging.getLogger("elasticsearch")
 
@@ -177,3 +180,6 @@ class Connection(object):
         raise HTTP_EXCEPTIONS.get(status_code, TransportError)(
             status_code, error_message, additional_info
         )
+
+    def _get_default_user_agent(self):
+        return "elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version())

elasticsearch/connection/http_requests.py

@@ -73,6 +73,7 @@ class RequestsHttpConnection(Connection):
         self.session = requests.Session()
         self.session.headers = headers or {}
         self.session.headers.setdefault("content-type", "application/json")
+        self.session.headers.setdefault("user-agent", self._get_default_user_agent())
         if http_auth is not None:
             if isinstance(http_auth, (tuple, list)):
                 http_auth = tuple(http_auth)

elasticsearch/connection/http_urllib3.py

@@ -127,6 +127,7 @@ class Urllib3HttpConnection(Connection):
             self.headers.update({"content-encoding": "gzip"})
 
         self.headers.setdefault("content-type", "application/json")
+        self.headers.setdefault("user-agent", self._get_default_user_agent())
         pool_class = urllib3.HTTPConnectionPool
         kw = {}
 

elasticsearch/transport.py

@@ -336,11 +336,6 @@ class Transport(object):
             ignore = params.pop("ignore", ())
             if isinstance(ignore, int):
                 ignore = (ignore,)
-
-        if headers is None:
-            headers = {}
-        headers["user-agent"] = "elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version())
-
         for attempt in range(self.max_retries + 1):
             connection = self.get_connection()
 

test_elasticsearch/test_connection.py

@@ -5,6 +5,7 @@ from mock import Mock, patch
 import urllib3
 import warnings
 from requests.auth import AuthBase
+from platform import python_version
 
 from elasticsearch.exceptions import (
     TransportError,
@@ -14,9 +15,9 @@ from elasticsearch.exceptions import (
 )
 from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection
 from elasticsearch.exceptions import ImproperlyConfigured
+from elasticsearch import __versionstr__
 from .test_cases import TestCase, SkipTest
 
-
 class TestUrllib3Connection(TestCase):
     def test_ssl_context(self):
         try:
@@ -48,6 +49,10 @@ class TestUrllib3Connection(TestCase):
         self.assertTrue(con.http_compress)
         self.assertEquals(con.headers["content-encoding"], "gzip")
 
+    def test_default_user_agent(self):
+        con = Urllib3HttpConnection()
+        self.assertEquals(con._get_default_user_agent(), "elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version()))
+
     def test_timeout_set(self):
         con = Urllib3HttpConnection(timeout=42)
         self.assertEquals(42, con.timeout)
@@ -55,7 +60,11 @@ class TestUrllib3Connection(TestCase):
     def test_keep_alive_is_on_by_default(self):
         con = Urllib3HttpConnection()
         self.assertEquals(
-            {"connection": "keep-alive", "content-type": "application/json"},
+            {
+                "connection": "keep-alive",
+                "content-type": "application/json",
+                "user-agent": con._get_default_user_agent(),
+            },
             con.headers,
         )
 
@@ -66,6 +75,7 @@ class TestUrllib3Connection(TestCase):
                 "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
                 "connection": "keep-alive",
                 "content-type": "application/json",
+                "user-agent": con._get_default_user_agent(),
             },
             con.headers,
         )
@@ -77,6 +87,7 @@ class TestUrllib3Connection(TestCase):
                 "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
                 "content-type": "application/json",
                 "connection": "keep-alive",
+                "user-agent": con._get_default_user_agent(),
             },
             con.headers,
         )
@@ -88,6 +99,7 @@ class TestUrllib3Connection(TestCase):
                 "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
                 "content-type": "application/json",
                 "connection": "keep-alive",
+                "user-agent": con._get_default_user_agent(),
             },
             con.headers,
         )
@@ -213,6 +225,21 @@ class TestRequestsConnection(TestCase):
         self.assertEquals(req.headers["h2"], "v2p")
         self.assertEquals(req.headers["h3"], "v3")
 
+    def test_default_headers(self):
+        con = self._get_mock_connection()
+        req = self._get_request(con, "GET", "/")
+        self.assertEquals(req.headers["content-type"], "application/json")
+        self.assertEquals(req.headers["user-agent"], con._get_default_user_agent())
+
+    def test_custom_headers(self):
+        con = self._get_mock_connection()
+        req = self._get_request(con, "GET", "/", headers={
+            "content-type": "application/x-ndjson",
+            "user-agent": "custom-agent/1.2.3",
+        })
+        self.assertEquals(req.headers["content-type"], "application/x-ndjson")
+        self.assertEquals(req.headers["user-agent"], "custom-agent/1.2.3")
+
     def test_http_auth(self):
         con = RequestsHttpConnection(http_auth="username:secret")
         self.assertEquals(("username", "secret"), con.session.auth)

test_elasticsearch/test_transport.py

@@ -1,13 +1,11 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 import time
-from platform import python_version
 
 from elasticsearch.transport import Transport, get_host_info
 from elasticsearch.connection import Connection
 from elasticsearch.connection_pool import DummyConnectionPool
 from elasticsearch.exceptions import ConnectionError, ImproperlyConfigured
-from elasticsearch import __versionstr__
 
 from .test_cases import TestCase
 
@@ -84,8 +82,21 @@ class TestTransport(TestCase):
         self.assertEquals(1, len(t.get_connection().calls))
         self.assertEquals(("GET", "/", {}, None), t.get_connection().calls[0][0])
         self.assertEquals(
-            {"timeout": 42, "ignore": (), "headers": {
-                'user-agent':"elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version())}
+            {
+                "timeout": 42,
+                "ignore": (),
+                "headers": None,
+            },
+            t.get_connection().calls[0][1],
+        )
+
+    def test_request_with_custom_user_agent_header(self):
+        t = Transport([{}], connection_class=DummyConnection)
+
+        t.perform_request("GET", "/", headers={"user-agent": "my-custom-value/1.2.3"})
+        self.assertEquals(1, len(t.get_connection().calls))
+        self.assertEquals(
+            {"timeout": None, "ignore": (), "headers": {"user-agent": "my-custom-value/1.2.3"}
             },
             t.get_connection().calls[0][1],
         )