Adding GZip support to urllib3 (elastic#704)

* Adding GZip support to urllib3 * Adding compression documentation and example * Convert to lowercase for consistency * Moving header manipulation to __init__() * Validating headers for compression * Moving body compression out of the headers block * Don't compress if there is no body * Infer true

Author: robgil

docs/index.rst

@@ -258,6 +258,17 @@ bodies via post::
  from elasticsearch import Elasticsearch
  es = Elasticsearch(send_get_body_as='POST')
 
+Compression
+~~~~~~~~~~~
+When using capacity constrained networks (low throughput), it may be handy to enable
+compression. This is especially useful when doing bulk loads or inserting large
+documents. This will configure compression on the *request*.
+::
+
+ from elasticsearch import Elasticsearch
+ es = Elasticsearch(hosts, http_compress = True)
+
+
 Running on AWS with IAM
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

elasticsearch/connection/http_urllib3.py

@@ -3,6 +3,7 @@
 import urllib3
 from urllib3.exceptions import ReadTimeoutError, SSLError as UrllibSSLError
 import warnings
+import gzip
 
 # sentinal value for `verify_certs`.
 # This is used to detect if a user is passing in a value for `verify_certs`
@@ -62,13 +63,15 @@ class Urllib3HttpConnection(Connection):
  host. See https://urllib3.readthedocs.io/en/1.4/pools.html#api for more
  information.
  :arg headers: any custom http headers to be add to requests
+ :arg http_compress: Use gzip compression
  """
  def __init__(self, host='localhost', port=9200, http_auth=None,
  use_ssl=False, verify_certs=VERIFY_CERTS_DEFAULT, ca_certs=None, client_cert=None,
  client_key=None, ssl_version=None, ssl_assert_hostname=None,
- ssl_assert_fingerprint=None, maxsize=10, headers=None, ssl_context=None, **kwargs):
+ ssl_assert_fingerprint=None, maxsize=10, headers=None, ssl_context=None, http_compress=False, **kwargs):
 
  super(Urllib3HttpConnection, self).__init__(host=host, port=port, use_ssl=use_ssl, **kwargs)
+ self.http_compress = http_compress
  self.headers = urllib3.make_headers(keep_alive=True)
  if http_auth is not None:
  if isinstance(http_auth, (tuple, list)):
@@ -80,6 +83,10 @@ def __init__(self, host='localhost', port=9200, http_auth=None,
  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'})
+
  self.headers.setdefault('content-type', 'application/json')
  pool_class = urllib3.HTTPConnectionPool
  kw = {}
@@ -154,6 +161,8 @@ def perform_request(self, method, url, params=None, body=None, timeout=None, ign
  if headers:
  request_headers = request_headers.copy()
  request_headers.update(headers)
+ if self.http_compress and body:
+ body = gzip.compress(body)
  response = self.pool.urlopen(method, url, body, retries=False, headers=request_headers, **kw)
  duration = time.time() - start
  raw_data = response.data.decode('utf-8')

test_elasticsearch/test_connection.py

@@ -32,6 +32,11 @@ def test_ssl_context(self):
  )
  self.assertTrue(con.use_ssl)
 
+ def test_http_compression(self):
+ con = Urllib3HttpConnection(http_compress=True)
+ self.assertTrue(con.http_compress)
+ self.assertEquals(con.headers['content-encoding'], 'gzip')
+
  def test_timeout_set(self):
  con = Urllib3HttpConnection(timeout=42)
  self.assertEquals(42, con.timeout)