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

[7.x] Refactor cloud_id, api_key, headers, and http_compress

Browse Source
This commit is contained in:
Seth Michael Larson
2020-03-04 13:07:07 -06:00
committed by GitHub
parent b0bc6418e3
commit 04403ffb1d
6 changed files with 158 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/index.rst
+1
View File
@@ -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
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
elasticsearch/connection/__init__.py
+7
View File
@@ -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",
]
elasticsearch/connection/base.py
+52 -9
View File
@@ -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
elasticsearch/connection/http_requests.py
+28 -28
View File
@@ -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
elasticsearch/connection/http_urllib3.py
+20 -32
View File
@@ -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(
test_elasticsearch/test_connection.py
+50
View File
@@ -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(