diff --git a/elasticsearch/__init__.py b/elasticsearch/__init__.py index 47aca528..590c69b7 100644 --- a/elasticsearch/__init__.py +++ b/elasticsearch/__init__.py @@ -72,7 +72,8 @@ try: raise ImportError from ._async.http_aiohttp import AIOHttpConnection + from ._async.transport import AsyncTransport - __all__ += ["AIOHttpConnection"] + __all__ += ["AIOHttpConnection", "AsyncTransport"] except (ImportError, SyntaxError): pass diff --git a/elasticsearch/_async/transport.py b/elasticsearch/_async/transport.py new file mode 100644 index 00000000..ac5d0b2b --- /dev/null +++ b/elasticsearch/_async/transport.py @@ -0,0 +1,331 @@ +# Licensed to Elasticsearch B.V under one or more agreements. +# Elasticsearch B.V licenses this file to you under the Apache 2.0 License. +# See the LICENSE file in the project root for more information + +import asyncio +import logging +from itertools import chain + +from .compat import get_running_loop +from .http_aiohttp import AIOHttpConnection +from ..transport import Transport +from ..exceptions import ( + TransportError, + ConnectionTimeout, + ConnectionError, + SerializationError, +) + + +logger = logging.getLogger("elasticsearch") + + +class AsyncTransport(Transport): + """ + Encapsulation of transport-related to logic. Handles instantiation of the + individual connections as well as creating a connection pool to hold them. + + Main interface is the `perform_request` method. + """ + + DEFAULT_CONNECTION_CLASS = AIOHttpConnection + + def __init__(self, hosts, *args, sniff_on_start=False, **kwargs): + """ + :arg hosts: list of dictionaries, each containing keyword arguments to + create a `connection_class` instance + :arg connection_class: subclass of :class:`~elasticsearch.Connection` to use + :arg connection_pool_class: subclass of :class:`~elasticsearch.ConnectionPool` to use + :arg host_info_callback: callback responsible for taking the node information from + `/_cluster/nodes`, along with already extracted information, and + producing a list of arguments (same as `hosts` parameter) + :arg sniff_on_start: flag indicating whether to obtain a list of nodes + from the cluster at startup time + :arg sniffer_timeout: number of seconds between automatic sniffs + :arg sniff_on_connection_fail: flag controlling if connection failure triggers a sniff + :arg sniff_timeout: timeout used for the sniff request - it should be a + fast api call and we are talking potentially to more nodes so we want + to fail quickly. Not used during initial sniffing (if + ``sniff_on_start`` is on) when the connection still isn't + initialized. + :arg serializer: serializer instance + :arg serializers: optional dict of serializer instances that will be + used for deserializing data coming from the server. (key is the mimetype) + :arg default_mimetype: when no mimetype is specified by the server + response assume this mimetype, defaults to `'application/json'` + :arg max_retries: maximum number of retries before an exception is propagated + :arg retry_on_status: set of HTTP status codes on which we should retry + on a different node. defaults to ``(502, 503, 504)`` + :arg retry_on_timeout: should timeout trigger a retry on different + node? (default `False`) + :arg send_get_body_as: for GET requests with body this option allows + you to specify an alternate way of execution for environments that + don't support passing bodies with GET requests. If you set this to + 'POST' a POST method will be used instead, if to 'source' then the body + will be serialized and passed as a query parameter `source`. + + Any extra keyword arguments will be passed to the `connection_class` + when creating and instance unless overridden by that connection's + options provided as part of the hosts parameter. + """ + self.sniffing_task = None + self.loop = None + self._async_init_called = False + + super(AsyncTransport, self).__init__( + *args, hosts=[], sniff_on_start=False, **kwargs + ) + + # Don't enable sniffing on Cloud instances. + if kwargs.get("cloud_id", False): + sniff_on_start = False + + # Since we defer connections / sniffing to not occur + # within the constructor we never want to signal to + # our parent to 'sniff_on_start' or non-empty 'hosts'. + self.hosts = hosts + self.sniff_on_start = sniff_on_start + + async def _async_init(self): + """This is our stand-in for an async constructor. Everything + that was deferred within __init__() should be done here now. + + This method will only be called once per AsyncTransport instance + and is called from one of AsyncElasticsearch.__aenter__(), + AsyncTransport.perform_request() or AsyncTransport.get_connection() + """ + # Detect the async loop we're running in and set it + # on all already created HTTP connections. + self.loop = get_running_loop() + self.kwargs["loop"] = self.loop + + # Now that we have a loop we can create all our HTTP connections + self.set_connections(self.hosts) + self.seed_connections = list(self.connection_pool.connections[:]) + + # ... and we can start sniffing in the background. + if self.sniffing_task is None and self.sniff_on_start: + self.last_sniff = self.loop.time() + self.create_sniff_task(initial=True) + + async def _async_call(self): + """This method is called within any async method of AsyncTransport + where the transport is not closing. This will check to see if we should + call our _async_init() or create a new sniffing task + """ + if not self._async_init_called: + self._async_init_called = True + await self._async_init() + + if self.sniffer_timeout: + if self.loop.time() >= self.last_sniff + self.sniff_timeout: + self.create_sniff_task() + + async def _get_node_info(self, conn, initial): + try: + # use small timeout for the sniffing request, should be a fast api call + _, headers, node_info = await conn.perform_request( + "GET", + "/_nodes/_all/http", + timeout=self.sniff_timeout if not initial else None, + ) + return self.deserializer.loads(node_info, headers.get("content-type")) + except Exception: + pass + return None + + async def _get_sniff_data(self, initial=False): + previous_sniff = self.last_sniff + + # reset last_sniff timestamp + self.last_sniff = self.loop.time() + + # use small timeout for the sniffing request, should be a fast api call + timeout = self.sniff_timeout if not initial else None + + def _sniff_request(conn): + return self.loop.create_task( + conn.perform_request("GET", "/_nodes/_all/http", timeout=timeout) + ) + + # Go through all current connections as well as the + # seed_connections for good measure + tasks = [] + for conn in self.connection_pool.connections: + tasks.append(_sniff_request(conn)) + for conn in self.seed_connections: + # Ensure that we don't have any duplication within seed_connections. + if conn in self.connection_pool.connections: + continue + tasks.append(_sniff_request(conn)) + + done = () + try: + while tasks: + # execute sniff requests in parallel, wait for first to return + done, tasks = await asyncio.wait( + tasks, return_when=asyncio.FIRST_COMPLETED, loop=self.loop + ) + # go through all the finished tasks + for t in done: + try: + _, headers, node_info = t.result() + node_info = self.deserializer.loads( + node_info, headers.get("content-type") + ) + except (ConnectionError, SerializationError): + continue + node_info = list(node_info["nodes"].values()) + return node_info + else: + # no task has finished completely + raise TransportError("N/A", "Unable to sniff hosts.") + except Exception: + # keep the previous value on error + self.last_sniff = previous_sniff + raise + finally: + # Cancel all the pending tasks + for task in chain(done, tasks): + task.cancel() + + async def sniff_hosts(self, initial=False): + """Either spawns a sniffing_task which does regular sniffing + over time or does a single sniffing session and awaits the results. + """ + # Without a loop we can't do anything. + if not self.loop: + return + + node_info = await self._get_sniff_data(initial) + hosts = list(filter(None, (self._get_host_info(n) for n in node_info))) + + # we weren't able to get any nodes, maybe using an incompatible + # transport_schema or host_info_callback blocked all - raise error. + if not hosts: + raise TransportError( + "N/A", "Unable to sniff hosts - no viable hosts found." + ) + + # remember current live connections + orig_connections = self.connection_pool.connections[:] + self.set_connections(hosts) + # close those connections that are not in use any more + for c in orig_connections: + if c not in self.connection_pool.connections: + await c.close() + + def create_sniff_task(self, initial=False): + """ + Initiate a sniffing task. Make sure we only have one sniff request + running at any given time. If a finished sniffing request is around, + collect its result (which can raise its exception). + """ + if self.sniffing_task and self.sniffing_task.done(): + try: + if self.sniffing_task is not None: + self.sniffing_task.result() + finally: + self.sniffing_task = None + + if self.sniffing_task is None: + self.sniffing_task = self.loop.create_task(self.sniff_hosts(initial)) + + def mark_dead(self, connection): + """ + Mark a connection as dead (failed) in the connection pool. If sniffing + on failure is enabled this will initiate the sniffing process. + + :arg connection: instance of :class:`~elasticsearch.Connection` that failed + """ + self.connection_pool.mark_dead(connection) + if self.sniff_on_connection_fail: + self.create_sniff_task() + + def get_connection(self): + return self.connection_pool.get_connection() + + async def perform_request(self, method, url, headers=None, params=None, body=None): + """ + Perform the actual request. Retrieve a connection from the connection + pool, pass all the information to it's perform_request method and + return the data. + + If an exception was raised, mark the connection as failed and retry (up + to `max_retries` times). + + If the operation was successful and the connection used was previously + marked as dead, mark it as live, resetting it's failure count. + + :arg method: HTTP method to use + :arg url: absolute url (without host) to target + :arg headers: dictionary of headers, will be handed over to the + underlying :class:`~elasticsearch.Connection` class + :arg params: dictionary of query parameters, will be handed over to the + underlying :class:`~elasticsearch.Connection` class for serialization + :arg body: body of the request, will be serialized using serializer and + passed to the connection + """ + await self._async_call() + + method, params, body, ignore, timeout = self._resolve_request_args( + method, params, body + ) + + for attempt in range(self.max_retries + 1): + connection = self.get_connection() + + try: + status, headers, data = await connection.perform_request( + method, + url, + params, + body, + headers=headers, + ignore=ignore, + timeout=timeout, + ) + except TransportError as e: + if method == "HEAD" and e.status_code == 404: + return False + + retry = False + if isinstance(e, ConnectionTimeout): + retry = self.retry_on_timeout + elif isinstance(e, ConnectionError): + retry = True + elif e.status_code in self.retry_on_status: + retry = True + + if retry: + # only mark as dead if we are retrying + self.mark_dead(connection) + # raise exception on last retry + if attempt == self.max_retries: + raise + else: + raise + + else: + if method == "HEAD": + return 200 <= status < 300 + + # connection didn't fail, confirm it's live status + self.connection_pool.mark_live(connection) + if data: + data = self.deserializer.loads(data, headers.get("content-type")) + return data + + async def close(self): + """ + Explicitly closes connections + """ + if self.sniffing_task: + try: + self.sniffing_task.cancel() + await self.sniffing_task + except asyncio.CancelledError: + pass + self.sniffing_task = None + for connection in self.connection_pool.connections: + await connection.close() diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py index 03f3949b..c0951c3e 100644 --- a/elasticsearch/connection_pool.py +++ b/elasticsearch/connection_pool.py @@ -256,9 +256,12 @@ class ConnectionPool(object): """ Explicitly closes connections """ - for conn in self.orig_connections: + for conn in self.connections: conn.close() + def __repr__(self): + return "<%s: %r>" % (type(self).__name__, self.connections) + class DummyConnectionPool(ConnectionPool): def __init__(self, connections, **kwargs): @@ -284,3 +287,19 @@ class DummyConnectionPool(ConnectionPool): pass mark_dead = mark_live = resurrect = _noop + + +class EmptyConnectionPool(ConnectionPool): + """A connection pool that is empty. Errors out if used.""" + + def __init__(self, *_, **__): + self.connections = [] + self.connection_opts = [] + + def get_connection(self): + raise ImproperlyConfigured("No connections were configured") + + def _noop(self, *args, **kwargs): + pass + + close = mark_dead = mark_live = resurrect = _noop diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py index faac81a5..79e8d540 100644 --- a/elasticsearch/transport.py +++ b/elasticsearch/transport.py @@ -6,7 +6,7 @@ import time from itertools import chain from .connection import Urllib3HttpConnection -from .connection_pool import ConnectionPool, DummyConnectionPool +from .connection_pool import ConnectionPool, DummyConnectionPool, EmptyConnectionPool from .serializer import JSONSerializer, Deserializer, DEFAULT_SERIALIZERS from .exceptions import ( ConnectionError, @@ -44,10 +44,12 @@ class Transport(object): Main interface is the `perform_request` method. """ + DEFAULT_CONNECTION_CLASS = Urllib3HttpConnection + def __init__( self, hosts, - connection_class=Urllib3HttpConnection, + connection_class=None, connection_pool_class=ConnectionPool, host_info_callback=get_host_info, sniff_on_start=False, @@ -100,6 +102,8 @@ class Transport(object): when creating and instance unless overridden by that connection's options provided as part of the hosts parameter. """ + if connection_class is None: + connection_class = self.DEFAULT_CONNECTION_CLASS # serialization config _serializers = DEFAULT_SERIALIZERS.copy() @@ -127,10 +131,18 @@ class Transport(object): self.kwargs = kwargs self.hosts = hosts - # ...and instantiate them - self.set_connections(hosts) - # retain the original connection instances for sniffing - self.seed_connections = self.connection_pool.connections[:] + # Start with an empty pool specifically for `AsyncTransport`. + # It should never be used, will be replaced on first call to + # .set_connections() + self.connection_pool = EmptyConnectionPool() + + if hosts: + # ...and instantiate them + self.set_connections(hosts) + # retain the original connection instances for sniffing + self.seed_connections = list(self.connection_pool.connections[:]) + else: + self.seed_connections = [] # Don't enable sniffing on Cloud instances. if kwargs.get("cloud_id", False): @@ -139,6 +151,7 @@ class Transport(object): # sniffing data self.sniffer_timeout = sniffer_timeout + self.sniff_on_start = sniff_on_start self.sniff_on_connection_fail = sniff_on_connection_fail self.last_sniff = time.time() self.sniff_timeout = sniff_timeout @@ -321,36 +334,9 @@ class Transport(object): :arg body: body of the request, will be serialized using serializer and passed to the connection """ - if body is not None: - body = self.serializer.dumps(body) - - # some clients or environments don't support sending GET with body - if method in ("HEAD", "GET") and self.send_get_body_as != "GET": - # send it as post instead - if self.send_get_body_as == "POST": - method = "POST" - - # or as source parameter - elif self.send_get_body_as == "source": - if params is None: - params = {} - params["source"] = body - body = None - - if body is not None: - try: - body = body.encode("utf-8", "surrogatepass") - except (UnicodeDecodeError, AttributeError): - # bytes/str - no need to re-encode - pass - - ignore = () - timeout = None - if params: - timeout = params.pop("request_timeout", None) - ignore = params.pop("ignore", ()) - if isinstance(ignore, int): - ignore = (ignore,) + method, params, body, ignore, timeout = self._resolve_request_args( + method, params, body + ) for attempt in range(self.max_retries + 1): connection = self.get_connection() @@ -405,3 +391,38 @@ class Transport(object): Explicitly closes connections """ self.connection_pool.close() + + def _resolve_request_args(self, method, params, body): + """Resolves parameters for .perform_request()""" + if body is not None: + body = self.serializer.dumps(body) + + # some clients or environments don't support sending GET with body + if method in ("HEAD", "GET") and self.send_get_body_as != "GET": + # send it as post instead + if self.send_get_body_as == "POST": + method = "POST" + + # or as source parameter + elif self.send_get_body_as == "source": + if params is None: + params = {} + params["source"] = body + body = None + + if body is not None: + try: + body = body.encode("utf-8", "surrogatepass") + except (UnicodeDecodeError, AttributeError): + # bytes/str - no need to re-encode + pass + + ignore = () + timeout = None + if params: + timeout = params.pop("request_timeout", None) + ignore = params.pop("ignore", ()) + if isinstance(ignore, int): + ignore = (ignore,) + + return method, params, body, ignore, timeout diff --git a/test_elasticsearch/test_async/test_transport.py b/test_elasticsearch/test_async/test_transport.py new file mode 100644 index 00000000..8e5b51a2 --- /dev/null +++ b/test_elasticsearch/test_async/test_transport.py @@ -0,0 +1,425 @@ +# -*- coding: utf-8 -*- +# Licensed to Elasticsearch B.V under one or more agreements. +# Elasticsearch B.V licenses this file to you under the Apache 2.0 License. +# See the LICENSE file in the project root for more information + +from __future__ import unicode_literals +import asyncio +from mock import patch +import pytest + +from elasticsearch import AsyncTransport +from elasticsearch.connection import Connection +from elasticsearch.connection_pool import DummyConnectionPool +from elasticsearch.exceptions import ConnectionError + + +pytestmark = pytest.mark.asyncio + + +class DummyConnection(Connection): + def __init__(self, **kwargs): + self.exception = kwargs.pop("exception", None) + self.status, self.data = kwargs.pop("status", 200), kwargs.pop("data", "{}") + self.headers = kwargs.pop("headers", {}) + self.delay = kwargs.pop("delay", 0) + self.calls = [] + self.closed = False + super(DummyConnection, self).__init__(**kwargs) + + async def perform_request(self, *args, **kwargs): + if self.closed: + raise RuntimeError("This connection is closed") + if self.delay: + await asyncio.sleep(self.delay) + self.calls.append((args, kwargs)) + if self.exception: + raise self.exception + return self.status, self.headers, self.data + + async def close(self): + if self.closed: + raise RuntimeError("This connection is already closed") + self.closed = True + + +CLUSTER_NODES = """{ + "_nodes" : { + "total" : 1, + "successful" : 1, + "failed" : 0 + }, + "cluster_name" : "elasticsearch", + "nodes" : { + "SRZpKFZdQguhhvifmN6UVA" : { + "name" : "SRZpKFZ", + "transport_address" : "127.0.0.1:9300", + "host" : "127.0.0.1", + "ip" : "127.0.0.1", + "version" : "5.0.0", + "build_hash" : "253032b", + "roles" : [ "master", "data", "ingest" ], + "http" : { + "bound_address" : [ "[fe80::1]:9200", "[::1]:9200", "127.0.0.1:9200" ], + "publish_address" : "1.1.1.1:123", + "max_content_length_in_bytes" : 104857600 + } + } + } +}""" + +CLUSTER_NODES_7x_PUBLISH_HOST = """{ + "_nodes" : { + "total" : 1, + "successful" : 1, + "failed" : 0 + }, + "cluster_name" : "elasticsearch", + "nodes" : { + "SRZpKFZdQguhhvifmN6UVA" : { + "name" : "SRZpKFZ", + "transport_address" : "127.0.0.1:9300", + "host" : "127.0.0.1", + "ip" : "127.0.0.1", + "version" : "5.0.0", + "build_hash" : "253032b", + "roles" : [ "master", "data", "ingest" ], + "http" : { + "bound_address" : [ "[fe80::1]:9200", "[::1]:9200", "127.0.0.1:9200" ], + "publish_address" : "somehost.tld/1.1.1.1:123", + "max_content_length_in_bytes" : 104857600 + } + } + } +}""" + + +class TestTransport: + async def test_single_connection_uses_dummy_connection_pool(self): + t = AsyncTransport([{}]) + await t._async_call() + assert isinstance(t.connection_pool, DummyConnectionPool) + t = AsyncTransport([{"host": "localhost"}]) + await t._async_call() + assert isinstance(t.connection_pool, DummyConnectionPool) + + async def test_request_timeout_extracted_from_params_and_passed(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + + await t.perform_request("GET", "/", params={"request_timeout": 42}) + assert 1 == len(t.get_connection().calls) + assert ("GET", "/", {}, None) == t.get_connection().calls[0][0] + assert { + "timeout": 42, + "ignore": (), + "headers": None, + } == t.get_connection().calls[0][1] + + async def test_opaque_id(self): + t = AsyncTransport([{}], opaque_id="app-1", connection_class=DummyConnection) + + await t.perform_request("GET", "/") + assert 1 == len(t.get_connection().calls) + assert ("GET", "/", None, None) == t.get_connection().calls[0][0] + assert { + "timeout": None, + "ignore": (), + "headers": None, + } == t.get_connection().calls[0][1] + + # Now try with an 'x-opaque-id' set on perform_request(). + await t.perform_request("GET", "/", headers={"x-opaque-id": "request-1"}) + assert 2 == len(t.get_connection().calls) + assert ("GET", "/", None, None) == t.get_connection().calls[1][0] + assert { + "timeout": None, + "ignore": (), + "headers": {"x-opaque-id": "request-1"}, + } == t.get_connection().calls[1][1] + + async def test_request_with_custom_user_agent_header(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + + await t.perform_request( + "GET", "/", headers={"user-agent": "my-custom-value/1.2.3"} + ) + assert 1 == len(t.get_connection().calls) + assert { + "timeout": None, + "ignore": (), + "headers": {"user-agent": "my-custom-value/1.2.3"}, + } == t.get_connection().calls[0][1] + + async def test_send_get_body_as_source(self): + t = AsyncTransport( + [{}], send_get_body_as="source", connection_class=DummyConnection + ) + + await t.perform_request("GET", "/", body={}) + assert 1 == len(t.get_connection().calls) + assert ("GET", "/", {"source": "{}"}, None) == t.get_connection().calls[0][0] + + async def test_send_get_body_as_post(self): + t = AsyncTransport( + [{}], send_get_body_as="POST", connection_class=DummyConnection + ) + + await t.perform_request("GET", "/", body={}) + assert 1 == len(t.get_connection().calls) + assert ("POST", "/", None, b"{}") == t.get_connection().calls[0][0] + + async def test_body_gets_encoded_into_bytes(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + + await t.perform_request("GET", "/", body="你好") + assert 1 == len(t.get_connection().calls) + assert ( + "GET", + "/", + None, + b"\xe4\xbd\xa0\xe5\xa5\xbd", + ) == t.get_connection().calls[0][0] + + async def test_body_bytes_get_passed_untouched(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + + body = b"\xe4\xbd\xa0\xe5\xa5\xbd" + await t.perform_request("GET", "/", body=body) + assert 1 == len(t.get_connection().calls) + assert ("GET", "/", None, body) == t.get_connection().calls[0][0] + + async def test_body_surrogates_replaced_encoded_into_bytes(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + + await t.perform_request("GET", "/", body="你好\uda6a") + assert 1 == len(t.get_connection().calls) + assert ( + "GET", + "/", + None, + b"\xe4\xbd\xa0\xe5\xa5\xbd\xed\xa9\xaa", + ) == t.get_connection().calls[0][0] + + async def test_kwargs_passed_on_to_connections(self): + t = AsyncTransport([{"host": "google.com"}], port=123) + await t._async_call() + assert 1 == len(t.connection_pool.connections) + assert "http://google.com:123" == t.connection_pool.connections[0].host + + async def test_kwargs_passed_on_to_connection_pool(self): + dt = object() + t = AsyncTransport([{}, {}], dead_timeout=dt) + await t._async_call() + assert dt is t.connection_pool.dead_timeout + + async def test_custom_connection_class(self): + class MyConnection(object): + def __init__(self, **kwargs): + self.kwargs = kwargs + + t = AsyncTransport([{}], connection_class=MyConnection) + await t._async_call() + assert 1 == len(t.connection_pool.connections) + assert isinstance(t.connection_pool.connections[0], MyConnection) + + def test_add_connection(self): + t = AsyncTransport([{}], randomize_hosts=False) + t.add_connection({"host": "google.com", "port": 1234}) + + assert 2 == len(t.connection_pool.connections) + assert "http://google.com:1234" == t.connection_pool.connections[1].host + + async def test_request_will_fail_after_X_retries(self): + t = AsyncTransport( + [{"exception": ConnectionError("abandon ship")}], + connection_class=DummyConnection, + ) + + connection_error = False + try: + await t.perform_request("GET", "/") + except ConnectionError: + connection_error = True + + assert connection_error + assert 4 == len(t.get_connection().calls) + + async def test_failed_connection_will_be_marked_as_dead(self): + t = AsyncTransport( + [{"exception": ConnectionError("abandon ship")}] * 2, + connection_class=DummyConnection, + ) + + connection_error = False + try: + await t.perform_request("GET", "/") + except ConnectionError: + connection_error = True + + assert connection_error + assert 0 == len(t.connection_pool.connections) + + async def test_resurrected_connection_will_be_marked_as_live_on_success(self): + t = AsyncTransport([{}, {}], connection_class=DummyConnection) + await t._async_call() + con1 = t.connection_pool.get_connection() + con2 = t.connection_pool.get_connection() + t.connection_pool.mark_dead(con1) + t.connection_pool.mark_dead(con2) + + await t.perform_request("GET", "/") + assert 1 == len(t.connection_pool.connections) + assert 1 == len(t.connection_pool.dead_count) + + async def test_sniff_will_use_seed_connections(self): + t = AsyncTransport([{"data": CLUSTER_NODES}], connection_class=DummyConnection) + await t._async_call() + t.set_connections([{"data": "invalid"}]) + + await t.sniff_hosts() + assert 1 == len(t.connection_pool.connections) + assert "http://1.1.1.1:123" == t.get_connection().host + + async def test_sniff_on_start_fetches_and_uses_nodes_list(self): + t = AsyncTransport( + [{"data": CLUSTER_NODES}], + connection_class=DummyConnection, + sniff_on_start=True, + ) + await t._async_call() + await t.sniffing_task # Need to wait for the sniffing task to complete + + assert 1 == len(t.connection_pool.connections) + assert "http://1.1.1.1:123" == t.get_connection().host + + async def test_sniff_on_start_ignores_sniff_timeout(self): + t = AsyncTransport( + [{"data": CLUSTER_NODES}], + connection_class=DummyConnection, + sniff_on_start=True, + sniff_timeout=12, + ) + await t._async_call() + await t.sniffing_task # Need to wait for the sniffing task to complete + + assert (("GET", "/_nodes/_all/http"), {"timeout": None}) == t.seed_connections[ + 0 + ].calls[0] + + async def test_sniff_uses_sniff_timeout(self): + t = AsyncTransport( + [{"data": CLUSTER_NODES}], + connection_class=DummyConnection, + sniff_timeout=42, + ) + await t._async_call() + await t.sniff_hosts() + + assert (("GET", "/_nodes/_all/http"), {"timeout": 42}) == t.seed_connections[ + 0 + ].calls[0] + + async def test_sniff_reuses_connection_instances_if_possible(self): + t = AsyncTransport( + [{"data": CLUSTER_NODES}, {"host": "1.1.1.1", "port": 123}], + connection_class=DummyConnection, + randomize_hosts=False, + ) + await t._async_call() + connection = t.connection_pool.connections[1] + connection.delay = 3.0 # Add this delay to make the sniffing deterministic. + + await t.sniff_hosts() + assert 1 == len(t.connection_pool.connections) + assert connection is t.get_connection() + + async def test_sniff_on_fail_triggers_sniffing_on_fail(self): + t = AsyncTransport( + [{"exception": ConnectionError("abandon ship")}, {"data": CLUSTER_NODES}], + connection_class=DummyConnection, + sniff_on_connection_fail=True, + max_retries=0, + randomize_hosts=False, + ) + await t._async_call() + + connection_error = False + try: + await t.perform_request("GET", "/") + except ConnectionError: + connection_error = True + + await t.sniffing_task # Need to wait for the sniffing task to complete + + assert connection_error + assert 1 == len(t.connection_pool.connections) + assert "http://1.1.1.1:123" == t.get_connection().host + + async def test_sniff_after_n_seconds(self, event_loop): + t = AsyncTransport( + [{"data": CLUSTER_NODES}], + connection_class=DummyConnection, + sniffer_timeout=5, + ) + await t._async_call() + + for _ in range(4): + await t.perform_request("GET", "/") + assert 1 == len(t.connection_pool.connections) + assert isinstance(t.get_connection(), DummyConnection) + t.last_sniff = event_loop.time() - 5.1 + + await t.perform_request("GET", "/") + await t.sniffing_task # Need to wait for the sniffing task to complete + + assert 1 == len(t.connection_pool.connections) + assert "http://1.1.1.1:123" == t.get_connection().host + assert event_loop.time() - 1 < t.last_sniff < event_loop.time() + 0.01 + + async def test_sniff_7x_publish_host(self): + # Test the response shaped when a 7.x node has publish_host set + # and the returend data is shaped in the fqdn/ip:port format. + t = AsyncTransport( + [{"data": CLUSTER_NODES_7x_PUBLISH_HOST}], + connection_class=DummyConnection, + sniff_timeout=42, + ) + await t._async_call() + await t.sniff_hosts() + # Ensure we parsed out the fqdn and port from the fqdn/ip:port string. + assert t.connection_pool.connection_opts[0][1] == { + "host": "somehost.tld", + "port": 123, + } + + @patch("elasticsearch._async.transport.AsyncTransport.sniff_hosts") + async def test_sniffing_disabled_on_cloud_instances(self, sniff_hosts): + t = AsyncTransport( + [{}], + sniff_on_start=True, + sniff_on_connection_fail=True, + connection_class=DummyConnection, + cloud_id="cluster:dXMtZWFzdC0xLmF3cy5mb3VuZC5pbyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5NyQ0ZmE4ODIxZTc1NjM0MDMyYmVkMWNmMjIxMTBlMmY5Ng==", + ) + await t._async_call() + + assert not t.sniff_on_connection_fail + assert sniff_hosts.call_args is None # Assert not called. + await t.perform_request("GET", "/", body={}) + assert 1 == len(t.get_connection().calls) + assert ("GET", "/", None, b"{}") == t.get_connection().calls[0][0] + + async def test_transport_close_closes_all_pool_connections(self): + t = AsyncTransport([{}], connection_class=DummyConnection) + await t._async_call() + + assert not any([conn.closed for conn in t.connection_pool.connections]) + await t.close() + assert all([conn.closed for conn in t.connection_pool.connections]) + + t = AsyncTransport([{}, {}], connection_class=DummyConnection) + await t._async_call() + + assert not any([conn.closed for conn in t.connection_pool.connections]) + await t.close() + assert all([conn.closed for conn in t.connection_pool.connections])