diff --git a/docs/index.rst b/docs/index.rst index f2d17d6d..9b1754c6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -296,6 +296,7 @@ documents. This will configure compression. from elasticsearch import Elasticsearch es = Elasticsearch(hosts, http_compress=True) +Compression is enabled by default when connecting to Elastic Cloud via ``cloud_id``. Running on AWS with IAM ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/elasticsearch/connection/__init__.py b/elasticsearch/connection/__init__.py index 2470bd39..e56e541d 100644 --- a/elasticsearch/connection/__init__.py +++ b/elasticsearch/connection/__init__.py @@ -1,3 +1,10 @@ from .base import Connection from .http_requests import RequestsHttpConnection from .http_urllib3 import Urllib3HttpConnection, create_ssl_context + +__all__ = [ + "Connection", + "RequestsHttpConnection", + "Urllib3HttpConnection", + "create_ssl_context", +] diff --git a/elasticsearch/connection/base.py b/elasticsearch/connection/base.py index a26945ce..e908fae5 100644 --- a/elasticsearch/connection/base.py +++ b/elasticsearch/connection/base.py @@ -1,5 +1,5 @@ import logging -import base64 +import binascii import gzip import io from platform import python_version @@ -9,7 +9,7 @@ try: except ImportError: import json -from ..exceptions import TransportError, HTTP_EXCEPTIONS +from ..exceptions import TransportError, ImproperlyConfigured, HTTP_EXCEPTIONS from .. import __versionstr__ logger = logging.getLogger("elasticsearch") @@ -29,6 +29,14 @@ class Connection(object): (`perform_request`) is thread-safe. Also responsible for logging. + + :arg host: hostname of the node (default: localhost) + :arg port: port to use (integer, default: 9200) + :arg use_ssl: use ssl for the connection if `True` + :arg url_prefix: optional url prefix for elasticsearch + :arg timeout: default timeout in seconds (float, default: 10) + :arg http_compress: Use gzip compression + :arg cloud_id: The Cloud ID from ElasticCloud. Convenient way to connect to cloud instances. """ def __init__( @@ -38,20 +46,55 @@ class Connection(object): use_ssl=False, url_prefix="", timeout=10, + headers=None, + http_compress=None, + cloud_id=None, + api_key=None, **kwargs ): - """ - :arg host: hostname of the node (default: localhost) - :arg port: port to use (integer, default: 9200) - :arg url_prefix: optional url prefix for elasticsearch - :arg timeout: default timeout in seconds (float, default: 10) - """ + + if cloud_id: + try: + _, cloud_id = cloud_id.split(":") + parent_dn, es_uuid, _ = ( + binascii.a2b_base64(cloud_id.encode("utf-8")).decode("utf-8").split("$") + ) + except ValueError: + raise ImproperlyConfigured("'cloud_id' is not properly formatted") + + host = "%s.%s" % (es_uuid, parent_dn) + port = 9243 + use_ssl = True + if http_compress is None: + http_compress = True + + # Work-around if the implementing class doesn't + # define the headers property before calling super().__init__() + if not hasattr(self, "headers"): + self.headers = {} + + headers = headers or {} + for key in headers: + self.headers[key.lower()] = headers[key] + + self.headers.setdefault("content-type", "application/json") + self.headers.setdefault("user-agent", self._get_default_user_agent()) + + if api_key is not None: + self.headers["authorization"] = self._get_api_key_header_val(api_key) + + if http_compress: + self.headers["accept-encoding"] = "gzip,deflate" + scheme = kwargs.get("scheme", "http") if use_ssl or scheme == "https": scheme = "https" use_ssl = True self.use_ssl = use_ssl + self.http_compress = http_compress or False + self.hostname = host + self.port = port self.host = "%s://%s:%s" % (scheme, host, port) if url_prefix: url_prefix = "/" + url_prefix.strip("/") @@ -200,5 +243,5 @@ class Connection(object): """ if isinstance(api_key, (tuple, list)): s = "{0}:{1}".format(api_key[0], api_key[1]).encode('utf-8') - return "ApiKey " + base64.b64encode(s).decode('utf-8') + return "ApiKey " + binascii.b2a_base64(s).rstrip(b"\r\n").decode('utf-8') return "ApiKey " + api_key diff --git a/elasticsearch/connection/http_requests.py b/elasticsearch/connection/http_requests.py index 43457fc6..733f694c 100644 --- a/elasticsearch/connection/http_requests.py +++ b/elasticsearch/connection/http_requests.py @@ -1,6 +1,5 @@ import time import warnings -from base64 import decodestring try: import requests @@ -35,9 +34,10 @@ class RequestsHttpConnection(Connection): :arg client_key: path to the file containing the private key if using separate cert and key files (client_cert will contain only the cert) :arg headers: any custom http headers to be add to requests - :arg cloud_id: The Cloud ID from ElasticCloud. Convient way to connect to cloud instances. - :arg api_key: optional API Key authentication as either base64 encoded string or a tuple. + :arg http_compress: Use gzip compression + :arg cloud_id: The Cloud ID from ElasticCloud. Convenient way to connect to cloud instances. Other host connection params will be ignored. + :arg api_key: optional API Key authentication as either base64 encoded string or a tuple. """ def __init__( @@ -52,7 +52,7 @@ class RequestsHttpConnection(Connection): client_cert=None, client_key=None, headers=None, - http_compress=False, + http_compress=None, cloud_id=None, api_key=None, **kwargs @@ -61,27 +61,24 @@ class RequestsHttpConnection(Connection): raise ImproperlyConfigured( "Please install requests to use RequestsHttpConnection." ) - if cloud_id: - cluster_name, cloud_id = cloud_id.split(":") - url, es_uuid, kibana_uuid = ( - decodestring(cloud_id.encode("utf-8")).decode("utf-8").split("$") - ) - 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: + # Initialize Session so .headers works before calling super().__init__(). + self.session = requests.Session() + for key in list(self.session.headers): + self.session.headers.pop(key) + + super(RequestsHttpConnection, self).__init__( + host=host, + port=port, + use_ssl=use_ssl, + headers=headers, + http_compress=http_compress, + cloud_id=cloud_id, + api_key=api_key, + **kwargs + ) + + if not self.http_compress: # Need to set this to 'None' otherwise Requests adds its own. self.session.headers["accept-encoding"] = None @@ -91,12 +88,11 @@ class RequestsHttpConnection(Connection): elif isinstance(http_auth, string_types): http_auth = tuple(http_auth.split(":", 1)) self.session.auth = http_auth - if api_key is not None: - self.session.headers['authorization'] = self._get_api_key_header_val(api_key) + self.base_url = "http%s://%s:%d%s" % ( "s" if self.use_ssl else "", - host, - port, + self.hostname, + self.port, self.url_prefix, ) self.session.verify = verify_certs @@ -189,6 +185,10 @@ class RequestsHttpConnection(Connection): return response.status_code, response.headers, raw_data + @property + def headers(self): + return self.session.headers + def close(self): """ Explicitly closes connections diff --git a/elasticsearch/connection/http_urllib3.py b/elasticsearch/connection/http_urllib3.py index ea174c7a..936d5d87 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 -from base64 import decodestring # sentinel value for `verify_certs` and `ssl_show_warn`. # This is used to detect if a user is passing in a value @@ -73,9 +72,9 @@ class Urllib3HttpConnection(Connection): information. :arg headers: any custom http headers to be add to requests :arg http_compress: Use gzip compression - :arg cloud_id: The Cloud ID from ElasticCloud. Convient way to connect to cloud instances. - :arg api_key: optional API Key authentication as either base64 encoded string or a tuple. + :arg cloud_id: The Cloud ID from ElasticCloud. Convenient way to connect to cloud instances. Other host connection params will be ignored. + :arg api_key: optional API Key authentication as either base64 encoded string or a tuple. """ def __init__( @@ -95,43 +94,29 @@ class Urllib3HttpConnection(Connection): maxsize=10, headers=None, ssl_context=None, - http_compress=False, + http_compress=None, cloud_id=None, api_key=None, **kwargs ): - - if cloud_id: - cluster_name, cloud_id = cloud_id.split(":") - url, es_uuid, kibana_uuid = ( - decodestring(cloud_id.encode("utf-8")).decode("utf-8").split("$") - ) - 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 - ) - self.http_compress = http_compress + # Initialize headers before calling super().__init__(). self.headers = urllib3.make_headers(keep_alive=True) + + super(Urllib3HttpConnection, self).__init__( + host=host, + port=port, + use_ssl=use_ssl, + headers=headers, + http_compress=http_compress, + cloud_id=cloud_id, + api_key=api_key, + **kwargs + ) if http_auth is not None: if isinstance(http_auth, (tuple, list)): http_auth = ":".join(http_auth) self.headers.update(urllib3.make_headers(basic_auth=http_auth)) - # update headers in lowercase to allow overriding of auth headers - if headers: - for k in headers: - self.headers[k.lower()] = headers[k] - - 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()) - if api_key is not None: - self.headers.setdefault('authorization', self._get_api_key_header_val(api_key)) pool_class = urllib3.HTTPConnectionPool kw = {} @@ -202,9 +187,12 @@ class Urllib3HttpConnection(Connection): if not ssl_show_warn: urllib3.disable_warnings() - self.pool = pool_class( - host, port=port, timeout=self.timeout, maxsize=maxsize, **kw + self.hostname, + port=self.port, + timeout=self.timeout, + maxsize=maxsize, + **kw ) def perform_request( diff --git a/test_elasticsearch/test_connection.py b/test_elasticsearch/test_connection.py index f084b7cd..fa1a94eb 100644 --- a/test_elasticsearch/test_connection.py +++ b/test_elasticsearch/test_connection.py @@ -65,6 +65,9 @@ class TestUrllib3Connection(TestCase): self.assertEquals( con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243" ) + self.assertEquals(con.port, 9243) + self.assertEquals(con.hostname, "0fd50f62320ed6539f6cb48e1b68.example.cloud.com") + self.assertTrue(con.http_compress) def test_api_key_auth(self): # test with tuple @@ -73,6 +76,7 @@ class TestUrllib3Connection(TestCase): api_key=("elastic", "changeme1"), ) self.assertEquals(con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTE=") + self.assertEquals(con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243") # test with base64 encoded string con = Urllib3HttpConnection( @@ -80,6 +84,7 @@ class TestUrllib3Connection(TestCase): api_key="ZWxhc3RpYzpjaGFuZ2VtZTI=", ) self.assertEquals(con.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=") + self.assertEquals(con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243") def test_no_http_compression(self): con = self._get_mock_connection() @@ -119,6 +124,26 @@ class TestUrllib3Connection(TestCase): self.assertEqual(kwargs["headers"]["accept-encoding"], "gzip,deflate") self.assertNotIn("content-encoding", kwargs["headers"]) + def test_cloud_id_http_compress_override(self): + # 'http_compress' will be 'True' by default for connections with + # 'cloud_id' set but should prioritize user-defined values. + con = Urllib3HttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + ) + self.assertEquals(con.http_compress, True) + + con = Urllib3HttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + http_compress=False + ) + self.assertEquals(con.http_compress, False) + + con = Urllib3HttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + http_compress=True + ) + self.assertEquals(con.http_compress, True) + def test_default_user_agent(self): con = Urllib3HttpConnection() self.assertEquals(con._get_default_user_agent(), "elasticsearch-py/%s (Python %s)" % (__versionstr__, python_version())) @@ -289,6 +314,9 @@ class TestRequestsConnection(TestCase): self.assertEquals( con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243" ) + self.assertEquals(con.port, 9243) + self.assertEquals(con.hostname, "0fd50f62320ed6539f6cb48e1b68.example.cloud.com") + self.assertTrue(con.http_compress) def test_api_key_auth(self): # test with tuple @@ -297,6 +325,7 @@ class TestRequestsConnection(TestCase): api_key=("elastic", "changeme1"), ) self.assertEquals(con.session.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTE=") + self.assertEquals(con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243") # test with base64 encoded string con = RequestsHttpConnection( @@ -304,6 +333,7 @@ class TestRequestsConnection(TestCase): api_key="ZWxhc3RpYzpjaGFuZ2VtZTI=", ) self.assertEquals(con.session.headers["authorization"], "ApiKey ZWxhc3RpYzpjaGFuZ2VtZTI=") + self.assertEquals(con.host, "https://0fd50f62320ed6539f6cb48e1b68.example.cloud.com:9243") def test_no_http_compression(self): con = self._get_mock_connection() @@ -340,6 +370,26 @@ class TestRequestsConnection(TestCase): self.assertNotIn("content-encoding", req.headers) self.assertEqual(req.headers["accept-encoding"], "gzip,deflate") + def test_cloud_id_http_compress_override(self): + # 'http_compress' will be 'True' by default for connections with + # 'cloud_id' set but should prioritize user-defined values. + con = RequestsHttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + ) + self.assertEquals(con.http_compress, True) + + con = RequestsHttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + http_compress=False + ) + self.assertEquals(con.http_compress, False) + + con = RequestsHttpConnection( + cloud_id="foobar:ZXhhbXBsZS5jbG91ZC5jb20kMGZkNTBmNjIzMjBlZDY1MzlmNmNiNDhlMWI2OCRhYzUzOTVhODgz\nNDU2NmM5ZjE1Y2Q4ZTQ5MGE=\n", + http_compress=True + ) + self.assertEquals(con.http_compress, True) + def test_uses_https_if_verify_certs_is_off(self): with warnings.catch_warnings(record=True) as w: con = self._get_mock_connection(