diff --git a/elasticsearch/connection/base.py b/elasticsearch/connection/base.py index a3fbd780..a26945ce 100644 --- a/elasticsearch/connection/base.py +++ b/elasticsearch/connection/base.py @@ -1,6 +1,7 @@ import logging import base64 - +import gzip +import io from platform import python_version try: @@ -70,6 +71,12 @@ class Connection(object): def __hash__(self): return id(self) + def _gzip_compress(self, body): + buf = io.BytesIO() + with gzip.GzipFile(fileobj=buf, mode="wb") as f: + f.write(body) + return buf.getvalue() + def _pretty_json(self, data): # pretty JSON in tracer curl logs try: diff --git a/elasticsearch/connection/http_requests.py b/elasticsearch/connection/http_requests.py index f9b6efd2..43457fc6 100644 --- a/elasticsearch/connection/http_requests.py +++ b/elasticsearch/connection/http_requests.py @@ -52,6 +52,7 @@ class RequestsHttpConnection(Connection): client_cert=None, client_key=None, headers=None, + http_compress=False, cloud_id=None, api_key=None, **kwargs @@ -68,14 +69,22 @@ class RequestsHttpConnection(Connection): host = "%s.%s" % (es_uuid, url) port = 9243 use_ssl = True - + http_compress = True super(RequestsHttpConnection, self).__init__( host=host, port=port, use_ssl=use_ssl, **kwargs ) + self.http_compress = http_compress 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 self.http_compress: + self.session.headers["accept-encoding"] = "gzip,deflate" + else: + # Need to set this to 'None' otherwise Requests adds its own. + self.session.headers["accept-encoding"] = None + if http_auth is not None: if isinstance(http_auth, (tuple, list)): http_auth = tuple(http_auth) @@ -116,9 +125,15 @@ class RequestsHttpConnection(Connection): self, method, url, params=None, body=None, timeout=None, ignore=(), headers=None ): url = self.base_url + url + headers = headers or {} if params: url = "%s?%s" % (url, urlencode(params or {})) + orig_body = body + if self.http_compress and body: + body = self._gzip_compress(body) + headers["content-encoding"] = "gzip" + start = time.time() request = requests.Request(method=method, headers=headers, url=url, data=body) prepared_request = self.session.prepare_request(request) @@ -155,7 +170,7 @@ class RequestsHttpConnection(Connection): method, url, response.request.path_url, - body, + orig_body, duration, response.status_code, raw_data, @@ -166,7 +181,7 @@ class RequestsHttpConnection(Connection): method, url, response.request.path_url, - body, + orig_body, response.status_code, raw_data, duration, diff --git a/elasticsearch/connection/http_urllib3.py b/elasticsearch/connection/http_urllib3.py index 6543e0e0..ea174c7a 100644 --- a/elasticsearch/connection/http_urllib3.py +++ b/elasticsearch/connection/http_urllib3.py @@ -4,7 +4,6 @@ import urllib3 from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError from urllib3.util.retry import Retry import warnings -import gzip from base64 import decodestring # sentinel value for `verify_certs` and `ssl_show_warn`. @@ -110,6 +109,7 @@ class Urllib3HttpConnection(Connection): host = "%s.%s" % (es_uuid, url) port = "9243" use_ssl = True + http_compress = True super(Urllib3HttpConnection, self).__init__( host=host, port=port, use_ssl=use_ssl, **kwargs ) @@ -125,9 +125,8 @@ class Urllib3HttpConnection(Connection): for k in headers: self.headers[k.lower()] = headers[k] - if self.http_compress == True: - self.headers.update(urllib3.make_headers(accept_encoding=True)) - self.headers.update({"content-encoding": "gzip"}) + if self.http_compress: + self.headers["accept-encoding"] = "gzip,deflate" self.headers.setdefault("content-type", "application/json") self.headers.setdefault("user-agent", self._get_default_user_agent()) @@ -217,6 +216,7 @@ class Urllib3HttpConnection(Connection): full_url = self.host + url start = time.time() + orig_body = body try: kw = {} if timeout: @@ -230,17 +230,12 @@ class Urllib3HttpConnection(Connection): if not isinstance(method, str): method = method.encode("utf-8") - request_headers = self.headers - if headers: - request_headers = request_headers.copy() - request_headers.update(headers) + request_headers = self.headers.copy() + request_headers.update(headers or ()) + if self.http_compress and body: - try: - body = gzip.compress(body) - except AttributeError: - # oops, Python2.7 doesn't have `gzip.compress` let's try - # again - body = gzip.zlib.compress(body) + body = self._gzip_compress(body) + request_headers["content-encoding"] = "gzip" response = self.pool.urlopen( method, url, body, retries=Retry(False), headers=request_headers, **kw @@ -249,7 +244,7 @@ class Urllib3HttpConnection(Connection): raw_data = response.data.decode("utf-8") except Exception as e: self.log_request_fail( - method, full_url, url, body, time.time() - start, exception=e + method, full_url, url, orig_body, time.time() - start, exception=e ) if isinstance(e, UrllibSSLError): raise SSLError("N/A", str(e), e) @@ -260,12 +255,12 @@ class Urllib3HttpConnection(Connection): # raise errors based on http status codes, let the client handle those if needed if not (200 <= response.status < 300) and response.status not in ignore: self.log_request_fail( - method, full_url, url, body, duration, response.status, raw_data + method, full_url, url, orig_body, duration, response.status, raw_data ) self._raise_error(response.status, raw_data) self.log_request_success( - method, full_url, url, body, response.status, raw_data, duration + method, full_url, url, orig_body, response.status, raw_data, duration ) return response.status, response.getheaders(), raw_data diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index 309041fd..f084b7cd 100644 --- a/test_elasticsearch/test_connection.py +++ b/test_elasticsearch/test_connection.py @@ -1,6 +1,7 @@ -import sys import re import ssl +import gzip +import io from mock import Mock, patch import urllib3 import warnings @@ -14,11 +15,32 @@ from elasticsearch.exceptions import ( NotFoundError, ) from elasticsearch.connection import RequestsHttpConnection, Urllib3HttpConnection -from elasticsearch.exceptions import ImproperlyConfigured from elasticsearch import __versionstr__ from .test_cases import TestCase, SkipTest + +def gzip_decompress(data): + buf = gzip.GzipFile(fileobj=io.BytesIO(data), mode="rb") + return buf.read() + + class TestUrllib3Connection(TestCase): + def _get_mock_connection( + self, connection_params={}, response_body=b"{}" + ): + con = Urllib3HttpConnection(**connection_params) + + def _dummy_urlopen(*args, **kwargs): + dummy_response = Mock() + dummy_response.headers = {} + dummy_response.status = 200 + dummy_response.data = response_body + _dummy_urlopen.call_args = (args, kwargs) + return dummy_response + + con.pool.urlopen = _dummy_urlopen + return con + def test_ssl_context(self): try: context = ssl.create_default_context() @@ -59,10 +81,43 @@ class TestUrllib3Connection(TestCase): ) self.assertEquals(con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=") + def test_no_http_compression(self): + con = self._get_mock_connection() + self.assertFalse(con.http_compress) + self.assertNotIn("accept-encoding", con.headers) + + con.perform_request("GET", "/") + + (_, _, req_body), kwargs = con.pool.urlopen.call_args + + self.assertFalse(req_body) + self.assertNotIn("accept-encoding", kwargs["headers"]) + self.assertNotIn("content-encoding", kwargs["headers"]) + def test_http_compression(self): - con = Urllib3HttpConnection(http_compress=True) + con = self._get_mock_connection({"http_compress": True}) self.assertTrue(con.http_compress) - self.assertEquals(con.headers["content-encoding"], "gzip") + self.assertEqual(con.headers["accept-encoding"], "gzip,deflate") + + # 'content-encoding' shouldn't be set at a connection level. + # Should be applied only if the request is sent with a body. + self.assertNotIn("content-encoding", con.headers) + + con.perform_request("GET", "/", body=b"{}") + + (_, _, req_body), kwargs = con.pool.urlopen.call_args + + self.assertEqual(gzip_decompress(req_body), b"{}") + self.assertEqual(kwargs["headers"]["accept-encoding"], "gzip,deflate") + self.assertEqual(kwargs["headers"]["content-encoding"], "gzip") + + con.perform_request("GET", "/") + + (_, _, req_body), kwargs = con.pool.urlopen.call_args + + self.assertFalse(req_body) + self.assertEqual(kwargs["headers"]["accept-encoding"], "gzip,deflate") + self.assertNotIn("content-encoding", kwargs["headers"]) def test_default_user_agent(self): con = Urllib3HttpConnection() @@ -171,6 +226,17 @@ class TestUrllib3Connection(TestCase): str(w[0].message), ) + @patch("elasticsearch.connection.base.logger") + def test_uncompressed_body_logged(self, logger): + con = self._get_mock_connection(connection_params={"http_compress": True}) + con.perform_request("GET", "/", body=b"{\"example\": \"body\"}") + + self.assertEquals(2, logger.debug.call_count) + req, resp = logger.debug.call_args_list + print(req, resp) + self.assertEquals('> {"example": "body"}', req[0][0] % req[0][1:]) + self.assertEquals('< {}', resp[0][0] % resp[0][1:]) + class TestRequestsConnection(TestCase): def _get_mock_connection( @@ -239,6 +305,41 @@ class TestRequestsConnection(TestCase): ) self.assertEquals(con.session.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=") + def test_no_http_compression(self): + con = self._get_mock_connection() + + self.assertFalse(con.http_compress) + self.assertNotIn("content-encoding", con.session.headers) + + con.perform_request("GET", "/") + + req = con.session.send.call_args[0][0] + self.assertNotIn("content-encoding", req.headers) + self.assertNotIn("accept-encoding", req.headers) + + def test_http_compression(self): + con = self._get_mock_connection( + {"http_compress": True}, + ) + + self.assertTrue(con.http_compress) + + # 'content-encoding' shouldn't be set at a session level. + # Should be applied only if the request is sent with a body. + self.assertNotIn("content-encoding", con.session.headers) + + con.perform_request("GET", "/", body=b"{}") + + req = con.session.send.call_args[0][0] + self.assertEqual(req.headers["content-encoding"], "gzip") + self.assertEqual(req.headers["accept-encoding"], "gzip,deflate") + + con.perform_request("GET", "/") + + req = con.session.send.call_args[0][0] + self.assertNotIn("content-encoding", req.headers) + self.assertEqual(req.headers["accept-encoding"], "gzip,deflate") + def test_uses_https_if_verify_certs_is_off(self): with warnings.catch_warnings(record=True) as w: con = self._get_mock_connection( @@ -400,6 +501,16 @@ class TestRequestsConnection(TestCase): self.assertEquals('> {"question": "what\'s that?"}', req[0][0] % req[0][1:]) self.assertEquals('< {"answer": "that\'s it!"}', resp[0][0] % resp[0][1:]) + @patch("elasticsearch.connection.base.logger") + def test_uncompressed_body_logged(self, logger): + con = self._get_mock_connection(connection_params={"http_compress": True}) + con.perform_request("GET", "/", body=b"{\"example\": \"body\"}") + + self.assertEquals(2, logger.debug.call_count) + req, resp = logger.debug.call_args_list + self.assertEquals('> {"example": "body"}', req[0][0] % req[0][1:]) + self.assertEquals('< {}', resp[0][0] % resp[0][1:]) + def test_defaults(self): con = self._get_mock_connection() request = self._get_request(con, "GET", "/")