diff --git a/esrally/client.py b/esrally/client.py index ff92b7553..4771c1fe9 100644 --- a/esrally/client.py +++ b/esrally/client.py @@ -20,6 +20,7 @@ import time import certifi +import elasticsearch import urllib3 from esrally import exceptions, doc_link @@ -127,18 +128,39 @@ def _is_set(self, client_opts, k): except KeyError: return False - def create(self): - import elasticsearch + def _postprocess_client_options(self): + """ + This method is used to do any non init processing of the client options, prior to creating a client + """ + if self._is_set(self.client_options, "use_api_key"): + self._generate_api_key() + + def _generate_api_key(self): + with self._create_sync_client() as instance: + # an api_key is generated per client, which happens after an initial handshake using an existing auth scheme + # once this key is generated, a new client is created using the api_key client option and http_auth is removed + # the reason the options were not removed in place was because of the logic to coalesce the key tuple into a + # header that the client could use. + api_key_response = instance.security.create_api_key({"name": "rally-api-key"}) + self.client_options.pop("http_auth", None) + self.client_options["api_key"] = (api_key_response["id"], api_key_response["api_key"]) + + def _create_sync_client(self): return elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) + def create(self): + self._postprocess_client_options() + return self._create_sync_client() + def create_async(self): - import elasticsearch import esrally.async_connection import io import aiohttp from elasticsearch.serializer import JSONSerializer + self._postprocess_client_options() + class LazyJSONSerializer(JSONSerializer): def loads(self, s): meta = RallyAsyncElasticsearch.request_context.get() diff --git a/tests/client_test.py b/tests/client_test.py index d1ebcf0d8..f9bfb3570 100644 --- a/tests/client_test.py +++ b/tests/client_test.py @@ -275,6 +275,24 @@ def test_create_https_connection_unverified_certificate_present_client_certifica self.assertDictEqual(original_client_options, client_options) + @mock.patch("elasticsearch.Elasticsearch") + def test_use_api_keys(self, _): + hosts = [{"host": "127.0.0.1", "port": 9200}] + client_options = {"use_api_key": True, "http_auth": ("a", "b")} + + f = client.EsClientFactory(hosts, client_options) + self.assertIsNotNone(f.client_options["http_auth"]) + f.create() + self.assertIsNone(f.client_options.get("http_auth", None)) + self.assertIsNotNone(f.client_options["api_key"]) + + # now assert that not having a http_auth client option does not cause issue + del f.client_options["api_key"] + self.assertIsNone(f.client_options.get("api_key", None)) + f.create() + self.assertIsNone(f.client_options.get("http_auth", None)) + self.assertIsNotNone(f.client_options["api_key"]) + class RestLayerTests(TestCase): @mock.patch("elasticsearch.Elasticsearch")