actions/opensearch-pyd
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[7.x] Add HTTP compression to RequestsHttpConnection

Browse Source
This commit is contained in:
Seth Michael Larson
2020-03-03 10:17:58 -06:00
committed by GitHub
parent 024abcfa39
commit b0bc6418e3
4 changed files with 153 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
elasticsearch/connection/base.py
+8 -1
View File
@@ -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
+18 -3
View File
@@ -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
+12 -17
View File
@@ -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
test_elasticsearch/test_connection.py
+115 -4
View File
@@ -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", "/")