Enabled api key elastic connections

Author: filipecosta90

engine/clients/elasticsearch/config.py

@@ -1,4 +1,39 @@
-ELASTIC_PORT = 9200
-ELASTIC_INDEX = "bench"
-ELASTIC_USER = "elastic"
-ELASTIC_PASSWORD = "passwd"
+import os
+from elasticsearch import Elasticsearch
+
+ELASTIC_PORT = int(os.getenv("ELASTIC_PORT", 9200))
+ELASTIC_INDEX = os.getenv("ELASTIC_INDEX", "bench")
+ELASTIC_USER = os.getenv("ELASTIC_USER", "elastic")
+ELASTIC_PASSWORD = os.getenv("ELASTIC_PASSWORD", "passwd")
+ELASTIC_API_KEY = os.getenv("ELASTIC_API_KEY", None)
+ELASTIC_TIMEOUT = int(os.getenv("ELASTIC_TIMEOUT", 90))
+
+
+def get_es_client(host, connection_params):
+    client: Elasticsearch = None
+    init_params = {
+        **{
+            "verify_certs": False,
+            "request_timeout": ELASTIC_TIMEOUT,
+            "retry_on_timeout": True,
+        },
+        **connection_params,
+    }
+    if host.startswith("http"):
+        url = ""
+    else:
+        url = "http://"
+    url += f"{host}:{ELASTIC_PORT}"
+    if ELASTIC_API_KEY is None:
+        client = Elasticsearch(
+            url,
+            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
+            **init_params,
+        )
+    else:
+        client = Elasticsearch(
+            url,
+            api_key=ELASTIC_API_KEY,
+            **init_params,
+        )
+    return client

engine/clients/elasticsearch/configure.py

@@ -6,9 +6,7 @@
 from engine.base_client.distances import Distance
 from engine.clients.elasticsearch.config import (
     ELASTIC_INDEX,
-    ELASTIC_PASSWORD,
-    ELASTIC_PORT,
-    ELASTIC_USER,
+    get_es_client,
 )
 
 
@@ -25,19 +23,7 @@ class ElasticConfigurator(BaseConfigurator):
 
     def __init__(self, host, collection_params: dict, connection_params: dict):
         super().__init__(host, collection_params, connection_params)
-        init_params = {
-            **{
-                "verify_certs": False,
-                "request_timeout": 90,
-                "retry_on_timeout": True,
-            },
-            **connection_params,
-        }
-        self.client = Elasticsearch(
-            f"http://{host}:{ELASTIC_PORT}",
-            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
-            **init_params,
-        )
+        self.client = get_es_client(host, connection_params)
 
     def clean(self):
         try:

engine/clients/elasticsearch/search.py

@@ -7,9 +7,7 @@
 from engine.base_client.search import BaseSearcher
 from engine.clients.elasticsearch.config import (
     ELASTIC_INDEX,
-    ELASTIC_PASSWORD,
-    ELASTIC_PORT,
-    ELASTIC_USER,
+    get_es_client,
 )
 from engine.clients.elasticsearch.parser import ElasticConditionParser
 
@@ -38,11 +36,7 @@ def init_client(cls, host, distance, connection_params: dict, search_params: dic
             },
             **connection_params,
         }
-        cls.client: Elasticsearch = Elasticsearch(
-            f"http://{host}:{ELASTIC_PORT}",
-            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
-            **init_params,
-        )
+        cls.client = get_es_client(host, connection_params)
         cls.search_params = search_params
 
     @classmethod

engine/clients/elasticsearch/upload.py

@@ -7,9 +7,7 @@
 from engine.base_client.upload import BaseUploader
 from engine.clients.elasticsearch.config import (
     ELASTIC_INDEX,
-    ELASTIC_PASSWORD,
-    ELASTIC_PORT,
-    ELASTIC_USER,
+    get_es_client,
 )
 
 
@@ -28,19 +26,7 @@ def get_mp_start_method(cls):
 
     @classmethod
     def init_client(cls, host, distance, connection_params, upload_params):
-        init_params = {
-            **{
-                "verify_certs": False,
-                "request_timeout": 90,
-                "retry_on_timeout": True,
-            },
-            **connection_params,
-        }
-        cls.client = Elasticsearch(
-            f"http://{host}:{ELASTIC_PORT}",
-            basic_auth=(ELASTIC_USER, ELASTIC_PASSWORD),
-            **init_params,
-        )
+        cls.client = get_es_client(host, connection_params)
         cls.upload_params = upload_params
 
     @classmethod