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

Author: philkra

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 @@ def _raise_error(self, status_code, raw_data):
  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 @@ def __init__(
  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 @@ def __init__(
  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 @@ def perform_request(self, method, url, headers=None, params=None, body=None):
  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 @@
 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.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,14 +49,22 @@ def test_http_compression(self):
  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)
 
  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 @@ def test_http_auth(self):
  "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
  "connection": "keep-alive",
  "content-type": "application/json",
+ "user-agent": con._get_default_user_agent(),
  },
  con.headers,
  )
@@ -77,6 +87,7 @@ def test_http_auth_tuple(self):
  "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
  "content-type": "application/json",
  "connection": "keep-alive",
+ "user-agent": con._get_default_user_agent(),
  },
  con.headers,
  )
@@ -88,6 +99,7 @@ def test_http_auth_list(self):
  "authorization": "Basic dXNlcm5hbWU6c2VjcmV0",
  "content-type": "application/json",
  "connection": "keep-alive",
+ "user-agent": con._get_default_user_agent(),
  },
  con.headers,
  )
@@ -213,6 +225,21 @@ def test_merge_headers(self):
  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 @@ def test_request_timeout_extracted_from_params_and_passed(self):
  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],
  )