Refactor elasticsearch logger into class

minecraft_discord_bridge/elasticsearch_logger.py

@@ -4,67 +4,55 @@ import logging
 
 import requests
 
-_username = None
-_password = None
-_auth = None
-_url = None
 
-log = logging.getLogger("bridge.elasticsearch")
-
-
-def initialize(config):
-    global _username, _password, _url, _auth
-    if config.es_auth:
-        _auth = True
-        _username = config.es_username
-        _password = config.es_password
-    _url = config.es_url
-
-
-def log_connection(uuid, reason, count=0):
-    if ConnectionReason(reason).name != "SEEN":
+class ElasticsearchLogger():
+    def __init__(self, url: str, username: str = "", password: str = ""):
+        self.url = url
+        self.username = username
+        self.password = password
+        self.log = logging.getLogger("bridge.elasticsearch")
+
+    def log_connection(self, uuid, reason, count=0):
+        if ConnectionReason(reason).name != "SEEN":
+            es_payload = {
+                "uuid": uuid,
+                "time": (lambda: int(round(time.time() * 1000)))(),
+                "reason": ConnectionReason(reason).name,
+                "count": count,
+            }
+        else:
+            es_payload = {
+                "uuid": uuid,
+                "time": (lambda: int(round(time.time() * 1000)))(),
+                "reason": ConnectionReason(reason).name,
+            }
+        self.post_request("connections/_doc/", es_payload)
+
+    def log_chat_message(self, uuid, display_name, message, message_unformatted):
         es_payload = {
             "uuid": uuid,
+            "display_name": display_name,
+            "message": message,
+            "message_unformatted": message_unformatted,
             "time": (lambda: int(round(time.time() * 1000)))(),
-            "reason": ConnectionReason(reason).name,
-            "count": count,
         }
-    else:
+        self.post_request("chat_messages/_doc/", es_payload)
+
+    def log_raw_message(self, type, message):
         es_payload = {
-            "uuid": uuid,
             "time": (lambda: int(round(time.time() * 1000)))(),
-            "reason": ConnectionReason(reason).name,
+            "type": type,
+            "message": message,
         }
-    post_request("connections/_doc/", es_payload)
-
-
-def log_chat_message(uuid, display_name, message, message_unformatted):
-    es_payload = {
-        "uuid": uuid,
-        "display_name": display_name,
-        "message": message,
-        "message_unformatted": message_unformatted,
-        "time": (lambda: int(round(time.time() * 1000)))(),
-    }
-    post_request("chat_messages/_doc/", es_payload)
-
-
-def log_raw_message(type, message):
-    es_payload = {
-        "time": (lambda: int(round(time.time() * 1000)))(),
-        "type": type,
-        "message": message,
-    }
-    post_request("raw_messages/_doc/", es_payload)
-
-
-def post_request(endpoint, payload):
-    the_url = "{}{}".format(_url, endpoint)
-    if _auth:
-        post = requests.post(the_url, auth=(_username, _password), json=payload)
-    else:
-        post = requests.post(the_url, json=payload)
-    log.debug(post.text)
+        self.post_request("raw_messages/_doc/", es_payload)
+
+    def post_request(self, endpoint, payload):
+        theURL = "{}{}".format(self.url, endpoint)
+        if self.username and self.password:
+            post = requests.post(theURL, auth=(self.username, self.password), json=payload)
+        else:
+            post = requests.post(theURL, json=payload)
+        self.log.debug(post.text)
 
 
 class ConnectionReason(Enum):