From 28920b5d595f937e1107d59550c45e7be53d3285 Mon Sep 17 00:00:00 2001 From: Philip Krauss <35487337+philkra@users.noreply.github.com> Date: Mon, 12 Aug 2019 14:20:10 +0200 Subject: [PATCH] Ensure a custom User-Agent header is not overwritten (#992) --- elasticsearch/connection/base.py | 6 +++++ elasticsearch/connection/http_requests.py | 1 + elasticsearch/connection/http_urllib3.py | 1 + elasticsearch/transport.py | 5 ---- test_elasticsearch/test_connection.py | 31 +++++++++++++++++++++-- test_elasticsearch/test_transport.py | 19 +++++++++++--- 6 files changed, 52 insertions(+), 11 deletions(-) diff --git a/elasticsearch/connection/base.py b/elasticsearch/connection/base.py index cb4a8116..45e338df 100644 --- a/elasticsearch/connection/base.py +++ b/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()) diff --git a/elasticsearch/connection/http_requests.py b/elasticsearch/connection/http_requests.py index e2a7a218..81282107 100644 --- a/elasticsearch/connection/http_requests.py +++ b/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) diff --git a/elasticsearch/connection/http_urllib3.py b/elasticsearch/connection/http_urllib3.py index d698f749..91078ad3 100644 --- a/elasticsearch/connection/http_urllib3.py +++ b/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 = {} diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py index 86b1c678..f8971b16 100644 --- a/elasticsearch/transport.py +++ b/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() diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index e7e9de72..1ee2e41f 100644 --- a/test_elasticsearch/test_connection.py +++ b/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) diff --git a/test_elasticsearch/test_transport.py b/test_elasticsearch/test_transport.py index d14306db..f06ed6be 100644 --- a/test_elasticsearch/test_transport.py +++ b/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], )