esrally/client/factory.py

# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you under
# the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# 	http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import contextlib
import logging
import time

import certifi
from urllib3.connection import is_ipaddress

from esrally import doc_link, exceptions
from esrally.utils import console, convert, versions


class EsClientFactory:
    """
    Abstracts how the Elasticsearch client is created and customizes the client for backwards
    compatibility guarantees that are broader than the library's defaults.
    """

    def __init__(self, hosts, client_options, distribution_version=None, distribution_flavor=None):
        def host_string(host):
            # protocol can be set at either host or client opts level
            protocol = "https" if client_options.get("use_ssl") or host.get("use_ssl") else "http"
            return f"{protocol}://{host['host']}:{host['port']}"

        self.hosts = [host_string(h) for h in hosts]
        self.client_options = dict(client_options)
        self.ssl_context = None
        # This attribute is necessary for the backwards-compatibility logic contained in
        # RallySyncElasticsearch.perform_request() and RallyAsyncElasticsearch.perform_request(), and also for
        # identification of whether or not a client is 'serverless'.
        self.distribution_version = distribution_version
        self.distribution_flavor = distribution_flavor
        self.logger = logging.getLogger(__name__)

        masked_client_options = dict(client_options)
        if "basic_auth_password" in masked_client_options:
            masked_client_options["basic_auth_password"] = "*****"
        if "http_auth" in masked_client_options:
            masked_client_options["http_auth"] = (masked_client_options["http_auth"][0], "*****")
        if "api_key" in masked_client_options:
            masked_client_options["api_key"] = "*****"
        self.logger.info("Creating ES client connected to %s with options [%s]", hosts, masked_client_options)

        # we're using an SSL context now and it is not allowed to have use_ssl present in client options anymore
        if self.client_options.pop("use_ssl", False):
            # pylint: disable=import-outside-toplevel
            import ssl

            self.logger.debug("SSL support: on")

            self.ssl_context = ssl.create_default_context(
                ssl.Purpose.SERVER_AUTH, cafile=self.client_options.pop("ca_certs", certifi.where())
            )

            # We call get() here instead of pop() in order to pass verify_certs through as a kwarg
            # to the elasticsearch.Elasticsearch constructor. Setting the ssl_context's verify_mode to
            # ssl.CERT_NONE is insufficient with version 8.0+ of elasticsearch-py.
            if not self.client_options.get("verify_certs", True):
                self.logger.debug("SSL certificate verification: off")
                # order matters to avoid ValueError: check_hostname needs a SSL context with either CERT_OPTIONAL or CERT_REQUIRED
                self.ssl_context.check_hostname = False
                self.ssl_context.verify_mode = ssl.CERT_NONE
                self.client_options["ssl_show_warn"] = False

                self.logger.debug(
                    "User has enabled SSL but disabled certificate verification. This is dangerous but may be ok for a benchmark."
                )
            else:
                # check_hostname should not be set when host is an IP address
                self.ssl_context.check_hostname = self._only_hostnames(hosts)
                self.ssl_context.verify_mode = ssl.CERT_REQUIRED
                self.logger.debug("SSL certificate verification: on")

            # When using SSL_context, all SSL related kwargs in client options get ignored
            client_cert = self.client_options.pop("client_cert", False)
            client_key = self.client_options.pop("client_key", False)

            if not client_cert and not client_key:
                self.logger.debug("SSL client authentication: off")
            elif bool(client_cert) != bool(client_key):
                self.logger.error("Supplied client-options contain only one of client_cert/client_key. ")
                defined_client_ssl_option = "client_key" if client_key else "client_cert"
                missing_client_ssl_option = "client_cert" if client_key else "client_key"
                console.println(
                    "'{}' is missing from client-options but '{}' has been specified.\n"
                    "If your Elasticsearch setup requires client certificate verification both need to be supplied.\n"
                    "Read the documentation at {}\n".format(
                        missing_client_ssl_option,
                        defined_client_ssl_option,
                        console.format.link(doc_link("command_line_reference.html#client-options")),
                    )
                )
                raise exceptions.SystemSetupError(
                    "Cannot specify '{}' without also specifying '{}' in client-options.".format(
                        defined_client_ssl_option, missing_client_ssl_option
                    )
                )
            elif client_cert and client_key:
                self.logger.debug("SSL client authentication: on")
                self.ssl_context.load_cert_chain(certfile=client_cert, keyfile=client_key)
        else:
            self.logger.debug("SSL support: off")

        if self._is_set(self.client_options, "create_api_key_per_client"):
            self.client_options.pop("create_api_key_per_client")
            basic_auth_user = self.client_options.get("basic_auth_user", False)
            basic_auth_password = self.client_options.get("basic_auth_password", False)
            provided_auth = {"basic_auth_user": basic_auth_user, "basic_auth_password": basic_auth_password}
            missing_auth = [k for k, v in provided_auth.items() if not v]
            if missing_auth:
                console.println(
                    "Basic auth credentials are required in order to create API keys.\n"
                    f"Missing basic auth client options are: {missing_auth}\n"
                    f"Read the documentation at {console.format.link(doc_link('command_line_reference.html#client-options'))}"
                )
                raise exceptions.SystemSetupError(
                    "You must provide the 'basic_auth_user' and 'basic_auth_password' client options in addition "
                    "to 'create_api_key_per_client' in order to create client API keys."
                )
            self.logger.debug("Automatic creation of client API keys: on")
        else:
            self.logger.debug("Automatic creation of client API keys: off")

        if self._is_set(self.client_options, "basic_auth_user") and self._is_set(self.client_options, "basic_auth_password"):
            self.client_options["basic_auth"] = (self.client_options.pop("basic_auth_user"), self.client_options.pop("basic_auth_password"))
            self.logger.debug("HTTP basic authentication: on")
        else:
            self.logger.debug("HTTP basic authentication: off")

        if self._is_set(self.client_options, "api_key"):
            self.logger.debug("API key authentication: on")
        else:
            self.logger.debug("API key authentication: off")

        if self._is_set(self.client_options, "compressed"):
            console.warn("You set the deprecated client option 'compressed‘. Please use 'http_compress' instead.", logger=self.logger)
            self.client_options["http_compress"] = self.client_options.pop("compressed")

        if self._is_set(self.client_options, "http_compress"):
            self.logger.debug("HTTP compression: on")
        else:
            self.logger.debug("HTTP compression: off")

        self.enable_cleanup_closed = convert.to_bool(self.client_options.pop("enable_cleanup_closed", True))
        self.max_connections = max(256, self.client_options.pop("max_connections", 0))
        self.static_responses = self.client_options.pop("static_responses", None)

        if self._is_set(self.client_options, "timeout"):
            self.client_options["request_timeout"] = self.client_options.pop("timeout")

    @staticmethod
    def _only_hostnames(hosts):
        has_ip = False
        has_hostname = False
        for host in hosts:
            is_ip = is_ipaddress(host["host"])
            if is_ip:
                has_ip = True
            else:
                has_hostname = True

        if has_ip and has_hostname:
            raise exceptions.SystemSetupError("Cannot verify certs with mixed IP addresses and hostnames")

        return has_hostname

    def _is_set(self, client_opts, k):
        try:
            return client_opts[k]
        except KeyError:
            return False

    def create(self):
        # pylint: disable=import-outside-toplevel
        from esrally.client.synchronous import RallySyncElasticsearch

        return RallySyncElasticsearch(
            distribution_version=self.distribution_version,
            distribution_flavor=self.distribution_flavor,
            hosts=self.hosts,
            ssl_context=self.ssl_context,
            **self.client_options,
        )

    def create_async(self, api_key=None, client_id=None):
        # pylint: disable=import-outside-toplevel
        import io

        import aiohttp
        from elasticsearch.serializer import JSONSerializer

        from esrally.client.asynchronous import (
            RallyAsyncElasticsearch,
            RallyAsyncTransport,
        )

        class LazyJSONSerializer(JSONSerializer):
            def loads(self, data):
                meta = RallyAsyncElasticsearch.request_context.get()
                if "raw_response" in meta:
                    return io.BytesIO(data)
                else:
                    return super().loads(data)

        async def on_request_start(session, trace_config_ctx, params):
            RallyAsyncElasticsearch.on_request_start()

        async def on_request_end(session, trace_config_ctx, params):
            RallyAsyncElasticsearch.on_request_end()

        trace_config = aiohttp.TraceConfig()
        trace_config.on_request_start.append(on_request_start)
        # It is tempting to register this callback on `TraceConfig.on_request_end()`. However, aiohttp will call
        # `TraceConfig.on_request_end()` when the *first* chunk of the response has been received. However, this can
        # skew service time significantly if the response is large *and* it is streamed by Elasticsearch
        # (see ChunkedToXContent in the Elasticsearch code base).
        #
        # Therefore, we register for `TraceConfig.on_response_chunk_received()` which is called multiple times. As
        # Rally's implementation of the `on_request_end` callback handler updates the timestamp on every call, Rally
        # will ultimately record the time when it received the *last* chunk. This is what we want because any code
        # that is using the Elasticsearch client library can only act on the response once it is fully received.
        #
        # We also keep registering for `TraceConfig.on_request_end()` instead of relying on
        # `TraceConfig.on_response_chunk_received()` only to handle corner cases during client timeout when aiohttp does
        # not call request exception handler, but does call request end handler. See
        # https://github.com/elastic/rally/issues/1860 for details.
        trace_config.on_response_chunk_received.append(on_request_end)
        trace_config.on_request_end.append(on_request_end)
        # ensure that we also stop the timer when a request "ends" with an exception (e.g. a timeout)
        trace_config.on_request_exception.append(on_request_end)

        # override the builtin JSON serializer
        self.client_options["serializer"] = LazyJSONSerializer()

        if api_key is not None:
            self.client_options.pop("http_auth", None)
            self.client_options.pop("basic_auth", None)
            self.client_options["api_key"] = api_key

        async_client = RallyAsyncElasticsearch(
            distribution_version=self.distribution_version,
            distribution_flavor=self.distribution_flavor,
            hosts=self.hosts,
            transport_class=RallyAsyncTransport,
            ssl_context=self.ssl_context,
            maxsize=self.max_connections,
            **self.client_options,
        )

        # the AsyncElasticsearch constructor automatically creates the corresponding NodeConfig objects, so we set
        # their instance attributes after they've been instantiated
        for node_connection in async_client.transport.node_pool.all():
            node_connection.trace_configs = [trace_config]
            node_connection.enable_cleanup_closed = self.enable_cleanup_closed
            node_connection.static_responses = self.static_responses
            node_connection.client_id = client_id

        return async_client


def wait_for_rest_layer(es, max_attempts=40):
    """
    Waits for ``max_attempts`` until Elasticsearch's REST API is available.

    :param es: Elasticsearch client to use for connecting.
    :param max_attempts: The maximum number of attempts to check whether the REST API is available.
    :return: True iff Elasticsearch's REST API is available.
    """
    # assume that at least the hosts that we expect to contact should be available. Note that this is not 100%
    # bullet-proof as a cluster could have e.g. dedicated masters which are not contained in our list of target hosts
    # but this is still better than just checking for any random node's REST API being reachable.
    expected_node_count = len(es.transport.node_pool)
    logger = logging.getLogger(__name__)
    attempt = 0
    while attempt <= max_attempts:
        attempt += 1
        # pylint: disable=import-outside-toplevel
        from elastic_transport import (
            ApiError,
            ConnectionError,
            SerializationError,
            TlsError,
            TransportError,
        )

        try:
            # see also WaitForHttpResource in Elasticsearch tests. Contrary to the ES tests we consider the API also
            # available when the cluster status is RED (as long as all required nodes are present)
            es.cluster.health(wait_for_nodes=f">={expected_node_count}")
            logger.debug("REST API is available for >= [%s] nodes after [%s] attempts.", expected_node_count, attempt)
            return True
        except SerializationError as e:
            if "Client sent an HTTP request to an HTTPS server" in str(e):
                raise exceptions.SystemSetupError(
                    "Rally sent an HTTP request to an HTTPS server. Are you sure this is an HTTP endpoint?", e
                )

            if attempt <= max_attempts:
                logger.debug("Got serialization error [%s] on attempt [%s]. Sleeping...", e, attempt)
                time.sleep(3)
            else:
                raise
        except TlsError as e:
            raise exceptions.SystemSetupError("Could not connect to cluster via HTTPS. Are you sure this is an HTTPS endpoint?", e)
        except ConnectionError as e:
            if "ProtocolError" in str(e):
                raise exceptions.SystemSetupError(
                    "Received a protocol error. Are you sure you're using the correct scheme (HTTP or HTTPS)?", e
                )

            if attempt <= max_attempts:
                logger.debug("Got connection error on attempt [%s]. Sleeping...", attempt)
                time.sleep(3)
            else:
                raise
        except TransportError as e:
            if attempt <= max_attempts:
                logger.debug("Got transport error on attempt [%s]. Sleeping...", attempt)
                time.sleep(3)
            else:
                raise
        except ApiError as e:
            # cluster block, x-pack not initialized yet, our wait condition is not reached
            if e.status_code in (503, 401, 408) and attempt <= max_attempts:
                logger.debug("Got status code [%s] on attempt [%s]. Sleeping...", e.message, attempt)
                time.sleep(3)
            else:
                logger.warning("Got unexpected status code [%s] on attempt [%s].", e.message, attempt)
                raise
    return False


def cluster_distribution_version(hosts, client_options, client_factory=EsClientFactory):
    """
    Attempt to get the target cluster's distribution version, build flavor, and build hash by creating and using
    a 'sync' Elasticsearch client.

    :param hosts: The host(s) to connect to.
    :param client_options: The client options to customize the Elasticsearch client.
    :param client_factory: Factory class that creates the Elasticsearch client.
    :return: The cluster's build flavor, version number, and build hash. For Serverless Elasticsearch these may all be
      the build flavor value. Also returns the operator status (always False for stateful).
    """
    # no way for us to know whether we're talking to a serverless elasticsearch or not, so we default to the sync client
    es = client_factory(hosts, client_options).create()

    # wait_for_rest_layer  calls the Cluster Health API, which is not available for unprivileged users on Serverless
    # As a result, we need to call the info API first to know if we can call wait_for_rest_layer().
    version = es.info()["version"]

    version_build_flavor = version.get("build_flavor", "oss")
    # if build hash is not available default to build flavor
    version_build_hash = version.get("build_hash", version_build_flavor)
    # if version number is not available default to build flavor
    version_number = version.get("number", version_build_flavor)

    # assume non-operator serverless privileges by default
    serverless_operator = False
    if versions.is_serverless(version_build_flavor):
        # overwrite static serverless version number
        version_number = "serverless"

        # determine serverless operator status if security enabled
        # pylint: disable=import-outside-toplevel
        from elasticsearch.exceptions import ApiError

        with contextlib.suppress(ApiError):
            authentication_info = es.perform_request(method="GET", path="/_security/_authenticate")
            serverless_operator = authentication_info.body.get("operator", False)

    if not versions.is_serverless(version_build_flavor) or serverless_operator is True:
        # if available, unconditionally wait for the REST layer - if it's not up, we'll intentionally raise the original error
        wait_for_rest_layer(es)

    return version_build_flavor, version_number, version_build_hash, serverless_operator


def create_api_key(es, client_id, max_attempts=5):
    """
    Creates an API key for the provided ``client_id``.

    :param es: Elasticsearch client to use for connecting.
    :param client_id: ID of the client for which the API key is being created.
    :param max_attempts: The maximum number of attempts to create the API key.
    :return: A dict with at least the following keys: ``id``, ``name``, ``api_key``.
    """
    logger = logging.getLogger(__name__)

    for attempt in range(1, max_attempts + 1):
        # pylint: disable=import-outside-toplevel
        import elasticsearch

        try:
            logger.debug("Creating ES API key for client ID [%s]", client_id)
            return es.security.create_api_key(name=f"rally-client-{client_id}")
        except elasticsearch.TransportError as e:
            logger.debug("Got transport error [%s] on attempt [%s]. Sleeping...", str(e), attempt)
            time.sleep(1)
        except elasticsearch.ApiError as e:
            if e.meta.status == 405:
                # We don't retry on 405 since it indicates a misconfigured benchmark candidate and isn't recoverable
                raise exceptions.SystemSetupError(
                    "Got status code 405 when attempting to create API keys. Is Elasticsearch Security enabled?", e
                )
            logger.debug("Got status code [%s] on attempt [%s] of [%s]. Sleeping...", e.status_code, attempt, max_attempts)
            time.sleep(1)


def delete_api_keys(es, ids, max_attempts=5):
    """
    Deletes the provided list of API key IDs.

    :param es: Elasticsearch client to use for connecting.
    :param ids: List of API key IDs to delete.
    :param max_attempts: The maximum number of attempts to delete the API keys.
    :return: True iff all provided key IDs were successfully deleted.
    """
    logger = logging.getLogger(__name__)

    def raise_exception(failed_ids, cause=None):
        msg = f"Could not delete API keys with the following IDs: {failed_ids}"
        if cause is not None:
            raise exceptions.RallyError(msg) from cause
        raise exceptions.RallyError(msg)

    # Before ES 7.10, deleting API keys by ID had to be done individually.
    # After ES 7.10, a list of API key IDs can be deleted in one request.
    version = es.info()["version"]
    current_version = versions.Version.from_string(version.get("number", "7.10.0"))
    minimum_version = versions.Version.from_string("7.10.0")

    deleted = []
    remaining = ids

    for attempt in range(1, max_attempts + 1):
        # pylint: disable=import-outside-toplevel
        import elasticsearch

        try:
            if current_version >= minimum_version or es.is_serverless:
                resp = es.security.invalidate_api_key(ids=remaining)
                deleted += resp["invalidated_api_keys"]
                remaining = [i for i in ids if i not in deleted]
                # Like bulk indexing requests, we can get an HTTP 200, but the
                # response body could still contain an array of individual errors.
                # So, we have to handle the case were some keys weren't deleted, but
                # the request overall succeeded (i.e. we didn't encounter an exception)
                if attempt < max_attempts:
                    if resp["error_count"] > 0:
                        logger.debug(
                            "Got the following errors on attempt [%s] of [%s]: [%s]. Sleeping...",
                            attempt,
                            max_attempts,
                            resp["error_details"],
                        )
                else:
                    if remaining:
                        logger.warning(
                            "Got the following errors on final attempt to delete API keys: [%s]",
                            resp["error_details"],
                        )
                        raise_exception(remaining)
            else:
                remaining = [i for i in ids if i not in deleted]
                if attempt < max_attempts:
                    for i in remaining:
                        es.security.invalidate_api_key(id=i)
                        deleted.append(i)
                else:
                    if remaining:
                        raise_exception(remaining)
            return True

        except elasticsearch.ApiError as e:
            if attempt < max_attempts:
                logger.debug("Got status code [%s] on attempt [%s] of [%s]. Sleeping...", e.meta.status, attempt, max_attempts)
                time.sleep(1)
            else:
                raise_exception(remaining, cause=e)
        except Exception as e:
            if attempt < max_attempts:
                logger.debug("Got error on attempt [%s] of [%s]. Sleeping...", attempt, max_attempts)
                time.sleep(1)
            else:
                raise_exception(remaining, cause=e)