From c37f3d8d5b9b8245a149952ef483faff99d2486f Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Mon, 15 Nov 2021 18:33:49 +0200
Subject: [PATCH] Fix set() not supported in ConfigTree() Add user/pass config
 support

---
 apiserver/es_factory.py | 64 ++++++++++++++++++++++++-----------------
 1 file changed, 38 insertions(+), 26 deletions(-)

diff --git a/apiserver/es_factory.py b/apiserver/es_factory.py
index cf3ca24..01e9516 100644
--- a/apiserver/es_factory.py
+++ b/apiserver/es_factory.py
@@ -1,9 +1,9 @@
 from datetime import datetime
 from os import getenv
-from typing import Tuple
+from typing import Tuple, Optional
 
 from boltons.iterutils import first
-from elasticsearch import Elasticsearch, Transport
+from elasticsearch import Elasticsearch
 
 from apiserver.config_repo import config
 
@@ -21,18 +21,18 @@ OVERRIDE_PORT_ENV_KEY = (
     "ELASTIC_SERVICE_PORT",
 )
 
-OVERRIDE_USERNAME_ENV_KEY = (
-    "CLEARML_ELASTIC_SERVICE_USERNAME",
-)
+OVERRIDE_USERNAME_ENV_KEY = ("CLEARML_ELASTIC_SERVICE_USERNAME",)
 
-OVERRIDE_PASSWORD_ENV_KEY = (
-    "CLEARML_ELASTIC_SERVICE_PASSWORD",
-)
+OVERRIDE_PASSWORD_ENV_KEY = ("CLEARML_ELASTIC_SERVICE_PASSWORD",)
 
 OVERRIDE_HOST = first(filter(None, map(getenv, OVERRIDE_HOST_ENV_KEY)))
 if OVERRIDE_HOST:
     log.info(f"Using override elastic host {OVERRIDE_HOST}")
 
+OVERRIDE_PORT = first(filter(None, map(getenv, OVERRIDE_PORT_ENV_KEY)))
+if OVERRIDE_PORT:
+    log.info(f"Using override elastic port {OVERRIDE_PORT}")
+
 OVERRIDE_USERNAME = first(filter(None, map(getenv, OVERRIDE_USERNAME_ENV_KEY)))
 if OVERRIDE_USERNAME:
     log.info(f"Using override elastic username {OVERRIDE_USERNAME}")
@@ -41,10 +41,6 @@ OVERRIDE_PASSWORD = first(filter(None, map(getenv, OVERRIDE_PASSWORD_ENV_KEY)))
 if OVERRIDE_PASSWORD:
     log.info("Using override elastic password ********")
 
-OVERRIDE_PORT = first(filter(None, map(getenv, OVERRIDE_PORT_ENV_KEY)))
-if OVERRIDE_PORT:
-    log.info(f"Using override elastic port {OVERRIDE_PORT}")
-
 _instances = {}
 
 
@@ -64,6 +60,10 @@ class InvalidClusterConfiguration(Exception):
     pass
 
 
+class MissingPasswordForElasticUser(Exception):
+    pass
+
+
 class ESFactory:
     @classmethod
     def connect(cls, cluster_name):
@@ -81,23 +81,42 @@ class ESFactory:
             if not hosts:
                 raise InvalidClusterConfiguration(cluster_name)
 
-            http_auth = cluster_config.get("http_auth", None)
-            if not cluster_config.get("secure", True):
-                http_auth = None
+            http_auth = (
+                cls.get_credentials(cluster_name)
+                if cluster_config.get("secure", True)
+                else None
+            )
+
             args = cluster_config.get("args", {})
             _instances[cluster_name] = Elasticsearch(
-                hosts=hosts, transport_class=Transport, http_auth=http_auth, **args
+                hosts=hosts, http_auth=http_auth, **args
             )
 
         return _instances[cluster_name]
 
+    @classmethod
+    def get_credentials(cls, cluster_name: str) -> Optional[Tuple[str, str]]:
+        elastic_user = OVERRIDE_USERNAME or config.get("secure.elastic.user", None)
+        if not elastic_user:
+            return None
+
+        elastic_password = OVERRIDE_PASSWORD or config.get(
+            "secure.elastic.password", None
+        )
+        if not elastic_password:
+            raise MissingPasswordForElasticUser(
+                f"cluster={cluster_name}, username={elastic_user}"
+            )
+
+        return elastic_user, elastic_password
+
     @classmethod
     def get_all_cluster_names(cls):
         return list(config.get("hosts.elastic"))
 
     @classmethod
-    def get_override(cls, cluster_name: str) -> Tuple[str, str, str, str]:
-        return OVERRIDE_HOST, OVERRIDE_PORT, OVERRIDE_USERNAME, OVERRIDE_PASSWORD
+    def get_override_host(cls, cluster_name: str) -> Tuple[str, str]:
+        return OVERRIDE_HOST, OVERRIDE_PORT
 
     @classmethod
     def get_cluster_config(cls, cluster_name):
@@ -116,7 +135,7 @@ class ESFactory:
             for entry in cluster_config.get("hosts", []):
                 entry[key] = value
 
-        host, port, username, password = cls.get_override(cluster_name)
+        host, port = cls.get_override_host(cluster_name)
 
         if host:
             set_host_prop("host", host)
@@ -124,13 +143,6 @@ class ESFactory:
         if port:
             set_host_prop("port", port)
 
-        if username and password:
-            cluster_config.set("http_auth", (username, password))
-        elif not username:
-            log.info(f"Ignoring username/password since username is empty or not provided")
-        elif not password:
-            log.info(f"Ignoring username/password since password is empty or not provided")
-
         return cluster_config
 
     @classmethod