Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optionally use generated api_key auth to ES #1098

Closed
wants to merge 2 commits into from
Closed

Optionally use generated api_key auth to ES #1098

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
61db7d8
Optionally use generated api_key auth to ES
hub-cap Oct 22, 2020
501f95a
Fixing up for review and adding a test
hub-cap Oct 26, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion esrally/client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would also enter the block if the user specified use_api_key:falseon the client options. I noticed that we fall into this trap in other places in this class as well and I propose we change this line to:

if self.client_options.get("use_api_key", False):

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That method is slightly misleading. It does not check if its set, it just removes the possibility of a key error. It returns the value, if present. Adding a test like so confirms this (and passes).

    def test_is_set(self):
        hosts = [{"host": "127.0.0.1", "port": 9200}]
        client_options = {}
        # make a copy so we can verify later that the factory did not modify it
        original_client_options = dict(client_options)

        f = client.EsClientFactory(hosts, client_options)

        opts = {"foo": True, "bar": False}
        assert f._is_set(opts, "foo") is True
        assert f._is_set(opts, "bar") is False
        assert f._is_set(opts, "baz") is False

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree that the semantics of the predicate is_set are misleading and maybe we should change it. I still propose we implement this line here differently:

if self.client_options.get("use_api_key", False):

This ensures correct behavior depending on the value of 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")
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know if it possible to use API keys without prior basic auth but pop raises a KeyError when the key is not present, so we should maybe do self.client_options.pop("http_auth", None) here, see docs.

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
Expand All @@ -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()
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that would work fine but how about we structure this a bit differently and instead encapsulate client option post processing into a dedicated method? Here is a summary of all my proposals (including earlier review comments:

diff --git a/esrally/client.py b/esrally/client.py
index ff92b75..ce76ccb 100644
--- a/esrally/client.py
+++ b/esrally/client.py
@@ -128,6 +128,22 @@ class EsClientFactory:
             return False

     def create(self):
+        self._postprocess_client_options()
+        return self._create_sync_client()
+
+    def _postprocess_client_options(self):
+        if self.client_options.get("use_api_key", False):
+            es = None
+            try:
+                es = self._create_sync_client()
+                api_key_response = es.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"])
+            finally:
+                if es:
+                    es.close()
+
+    def _create_sync_client(self):
         import elasticsearch
         return elasticsearch.Elasticsearch(hosts=self.hosts, ssl_context=self.ssl_context, **self.client_options)

@@ -163,6 +179,7 @@ class EsClientFactory:
         # 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)

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

Wdyt?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like this approach better 👍


class LazyJSONSerializer(JSONSerializer):
def loads(self, s):
meta = RallyAsyncElasticsearch.request_context.get()
Expand Down