[7.x] Add HTTP compression to RequestsHttpConnection

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:

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,

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:
+        if self.http_compress:
-            self.headers.update(urllib3.make_headers(accept_encoding=True))
+            self.headers["accept-encoding"] = "gzip,deflate"
-            self.headers.update({"content-encoding": "gzip"})
 
         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
+            request_headers = self.headers.copy()
-            if headers:
+            request_headers.update(headers or ())
-                request_headers = request_headers.copy()
+
-                request_headers.update(headers)
             if self.http_compress and body:
-                try:
+                body = self._gzip_compress(body)
-                    body = gzip.compress(body)
+                request_headers["content-encoding"] = "gzip"
-                except AttributeError:
-                    # oops, Python2.7 doesn't have `gzip.compress` let's try
-                    # again
-                    body = gzip.zlib.compress(body)
 
             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

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", "/")