Refactor of connection wrapper to handle errors connecting to ES properly.

tubearchivist-metrics/esconnect.py

@@ -1,15 +1,31 @@
-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
-    """
+    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
+        """
+        config = AppConfig().config
         es_url = config["es_url"]
         es_user = config["es_user"]
         es_pass = config["es_pass"]
@@ -17,14 +33,27 @@ class ElasticWrapper:
         es = Elasticsearch(
             [es_url],
             basic_auth=(es_user, es_pass),
-        timeout=30,
-        max_retries=10,
+            timeout=10,
+            max_retries=12,
             retry_on_timeout=True,
         )
 
-    def get_count(index_name):
-        """
-        Returns the number of documents in the index
-        """
-        response = ElasticWrapper.es.count(index=index_name)
-        return response["count"]
+        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