From d05f9aec5141242301a4425a6974d5a1f00097a2 Mon Sep 17 00:00:00 2001
From: Ainsey11 <robert@ainsey11.com>
Date: Fri, 15 Apr 2022 01:13:55 +0100
Subject: [PATCH] Refactor of connection wrapper to handle errors connecting to
 ES properly.

---
 tubearchivist-metrics/esconnect.py | 71 +++++++++++++++++++++---------
 1 file changed, 50 insertions(+), 21 deletions(-)

diff --git a/tubearchivist-metrics/esconnect.py b/tubearchivist-metrics/esconnect.py
index 094dc35..07abfe5 100644
--- a/tubearchivist-metrics/esconnect.py
+++ b/tubearchivist-metrics/esconnect.py
@@ -1,30 +1,59 @@
-from elasticsearch import Elasticsearch
-from environment import AppConfig
+from multiprocessing import AuthenticationError
+from elasticsearch import (
+    Elasticsearch,
+    ConnectionError,
+    ConnectionTimeout,
+    AuthenticationException,
+    AuthorizationException,
+)
 
-config = AppConfig().config
+from environment import AppConfig
+from time import sleep
 
 
 class ElasticWrapper:
-    """
-    makes calls to elastic search
-    returns response count
-    """
-
-    es_url = config["es_url"]
-    es_user = config["es_user"]
-    es_pass = config["es_pass"]
-
-    es = Elasticsearch(
-        [es_url],
-        basic_auth=(es_user, es_pass),
-        timeout=30,
-        max_retries=10,
-        retry_on_timeout=True,
-    )
+    def handle_err(error):
+        print("Connection Error: " + str(error))
+        print("There was a problem connecting to Elasticsearch")
+        print(
+            "Please see the error above. This may be as Elasticsearch is still starting up or a misconfiguration"
+        )
+        print("Sleeping for 10 seconds...")
+        sleep(10)
 
     def get_count(index_name):
         """
         Returns the number of documents in the index
         """
-        response = ElasticWrapper.es.count(index=index_name)
-        return response["count"]
+        config = AppConfig().config
+        es_url = config["es_url"]
+        es_user = config["es_user"]
+        es_pass = config["es_pass"]
+
+        es = Elasticsearch(
+            [es_url],
+            basic_auth=(es_user, es_pass),
+            timeout=10,
+            max_retries=12,
+            retry_on_timeout=True,
+        )
+
+        response = 0
+        try:
+            response = es.count(index=index_name)["count"]
+        except AuthenticationException as e:
+            ElasticWrapper.handle_err(e)
+        except ConnectionError as e:
+            ElasticWrapper.handle_err(e)
+        except ConnectionTimeout as e:
+            ElasticWrapper.handle_err(e)
+        except AuthenticationError as e:
+            ElasticWrapper.handle_err(e)
+        except AuthorizationException as e:
+            ElasticWrapper.handle_err(e)
+        except:
+            print("Unknown error occurred. Check your credentials, url and try again.")
+            print("Sleeping for 10 seconds...")
+            sleep(10)
+        else:
+            return response