From 64c125d34cdbb36ef29d2bf7dd587abeec1d9f5e Mon Sep 17 00:00:00 2001 From: robgil Date: Tue, 13 Mar 2018 12:34:25 -0400 Subject: [PATCH] Adding GZip support to urllib3 (#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 --- docs/index.rst | 11 +++++++++++ elasticsearch/connection/http_urllib3.py | 11 ++++++++++- test_elasticsearch/test_connection.py | 5 +++++ 3 files changed, 26 insertions(+), 1 deletion(-) diff --git a/docs/index.rst b/docs/index.rst index 0354a504..d196dad3 100644 --- a/docs/index.rst +++ b/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 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/elasticsearch/connection/http_urllib3.py b/elasticsearch/connection/http_urllib3.py index 36787d5a..9354307c 100644 --- a/elasticsearch/connection/http_urllib3.py +++ b/elasticsearch/connection/http_urllib3.py @@ -3,6 +3,7 @@ import ssl 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 @@ 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'}) + self.headers.setdefault('content-type', 'application/json') pool_class = urllib3.HTTPConnectionPool kw = {} @@ -154,6 +161,8 @@ class Urllib3HttpConnection(Connection): 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') diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index e03fdf63..430d43cb 100644 --- a/test_elasticsearch/test_connection.py +++ b/test_elasticsearch/test_connection.py @@ -32,6 +32,11 @@ class TestUrllib3Connection(TestCase): ) 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)