From 61db7d8c443757d109b20565e1719ce92da6105a Mon Sep 17 00:00:00 2001 From: Michael Basnight Date: Thu, 22 Oct 2020 16:11:47 -0500 Subject: [PATCH 1/2] Optionally use generated api_key auth to ES This commit adds a optional client arg, use_api_key, which will generate and use an api key for every client and async client created by Rally. It uses an initial call via the existing auth scheme to ES to generate the key, and then creates a new client using the api_key for general use by Rally. The async client also does this, and temporarily creates a sync client to generate the key, and then uses it when creating an async client. Closes #1067 --- esrally/client.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/esrally/client.py b/esrally/client.py index ff92b7553..052205588 100644 --- a/esrally/client.py +++ b/esrally/client.py @@ -129,7 +129,17 @@ def _is_set(self, client_opts, k): def create(self): import elasticsearch - return elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) + instance = elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) + if self._is_set(self.client_options, "use_api_key"): + # 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") + self.client_options["api_key"] = (api_key_response["id"], api_key_response["api_key"]) + instance = elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) + return instance def create_async(self): import elasticsearch @@ -139,6 +149,10 @@ def create_async(self): from elasticsearch.serializer import JSONSerializer + if self._is_set(self.client_options, "use_api_key"): + # Temporarily create a non async version of the client to generate an api_key and put it into client_options + self.create() + class LazyJSONSerializer(JSONSerializer): def loads(self, s): meta = RallyAsyncElasticsearch.request_context.get() From 501f95a73c9309eda08593b59318638f34d408bd Mon Sep 17 00:00:00 2001 From: Michael Basnight Date: Mon, 26 Oct 2020 14:46:13 -0500 Subject: [PATCH 2/2] Fixing up for review and adding a test --- esrally/client.py | 28 ++++++++++++++++++---------- tests/client_test.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+), 10 deletions(-) diff --git a/esrally/client.py b/esrally/client.py index 052205588..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,31 +128,38 @@ def _is_set(self, client_opts, k): except KeyError: return False - def create(self): - import elasticsearch - instance = elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) + 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") + self.client_options.pop("http_auth", None) self.client_options["api_key"] = (api_key_response["id"], api_key_response["api_key"]) - instance = elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options) - return instance + + 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 - if self._is_set(self.client_options, "use_api_key"): - # Temporarily create a non async version of the client to generate an api_key and put it into client_options - self.create() + self._postprocess_client_options() class LazyJSONSerializer(JSONSerializer): def loads(self, s): 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")