Added retry mechanism to InfluxDB connection creation.

README.md

@@ -12,13 +12,18 @@ This tool is a wrapper for speedtest-cli which allows you to run periodic speedt
 |:--------------|:-------------------------------------------------------------------------------------------------------------------|
 |Delay          |Delay between runs                                                                                                  |
 #### INFLUXDB
-|Key            |Description                                                                                                         |
-|:--------------|:-------------------------------------------------------------------------------------------------------------------|
-|Address        |Delay between updating metrics                                                                                      |
-|Port           |InfluxDB port to connect to.  8086 in most cases                                                                    |
-|Database       |Database to write collected stats to                                                                                |
-|Username       |User that has access to the database                                                                                |
-|Password       |Password for above user                                                                                             |
+| Key         | Description                                                |
+|:------------|:-----------------------------------------------------------|
+| Address     | Delay between updating metrics                             |
+| Port        | InfluxDB port to connect to.  8086 in most cases           |
+| Database    | Database to write collected stats to                       |
+| Username    | User that has access to the database                       |
+| Password    | Password for above user                                    |
+| Retry_Delay | Delay, in seconds, between attempts to connect to InfluxDB |
+| Retry_Count | Number of InfluxDB connection retries to attempt           |
+| Backoff     | Backoff multiplier                                         |
+
+
 #### SPEEDTEST
 |Key            |Description                                                                                                         |
 |:--------------|:-------------------------------------------------------------------------------------------------------------------|

config.ini

@@ -8,11 +8,14 @@ Database = speedtests
 Username =
 Password =
 Verify_SSL = False
+Retry_Delay = 5
+Retry_Count = 5
+Backoff = 1
 
 [SPEEDTEST]
 # Leave blank to auto pick server
 Server =
 
 [LOGGING]
 # Valid Options: critical, error, warning, info, debug
-Level = info
+Level = info

influxspeedtest/InfluxdbSpeedtest.py

@@ -7,17 +7,19 @@
 from requests import ConnectTimeout, ConnectionError
 
 from influxspeedtest.common import log
+from influxspeedtest.common.utils import retry
 from influxspeedtest.config import config
 
 
-class InfluxdbSpeedtest():
+class InfluxdbSpeedtest:
 
     def __init__(self):
 
         self.influx_client = self._get_influx_connection()
         self.speedtest = None
         self.results = None
 
+    @retry(exception=ConnectionError, n_tries=config.influx_retry_count, delay=config.influx_retry_delay, backoff=config.influx_retry_backoff)
     def _get_influx_connection(self):
         """
         Create an InfluxDB connection and test to make sure it works.
@@ -43,12 +45,11 @@ def _get_influx_connection(self):
         except (ConnectTimeout, InfluxDBClientError, ConnectionError) as e:
             if isinstance(e, ConnectTimeout):
                 log.critical('Unable to connect to InfluxDB at the provided address (%s)', config.influx_address)
-            elif e.code == 401:
-                log.critical('Unable to connect to InfluxDB with provided credentials')
+            elif isinstance(e, InfluxDBClientError):
+                log.critical('Unable to connect to InfluxDB with error code: %s', e.code)
             else:
                 log.critical('Failed to connect to InfluxDB for unknown reason')
-
-            sys.exit(1)
+            raise e
 
         return influx
 
@@ -65,7 +66,7 @@ def setup_speedtest(self, server=None):
         if server is None:
             server = []
         else:
-            server = server.split() # Single server to list
+            server = server.split()  # Single server to list
 
         try:
             self.speedtest = speedtest.Speedtest()
@@ -140,8 +141,6 @@ def run_speed_test(self, server=None):
                  results['server']['latency']
                  )
 
-
-
     def write_influx_data(self, json_data):
         """
         Writes the provided JSON to the database

influxspeedtest/common/utils.py

@@ -1,8 +1,10 @@
 import logging
 import sys
+import time
 
 from influxspeedtest.common.logfilters import SingleLevelFilter
 from influxspeedtest.config import config
+from functools import partial, wraps
 
 log = logging.getLogger(__name__)
 log.setLevel(config.logging_level)
@@ -20,4 +22,73 @@
 error_handler.addFilter(error_filter)
 log.addHandler(error_handler)
 
-log.propagate = False
+log.propagate = False
+
+
+def retry(func=None, exception=Exception, n_tries=5, delay=5, backoff=1, logger=True):
+    """Retry decorator with exponential backoff.
+
+    Parameters
+    ----------
+    func : typing.Callable, optional
+        Callable on which the decorator is applied, by default None
+    exception : Exception or tuple of Exceptions, optional
+        Exception(s) that invoke retry, by default Exception
+    n_tries : int, optional
+        Number of tries before giving up, by default 5
+    delay : int, optional
+        Initial delay between retries in seconds, by default 5
+    backoff : int, optional
+        Backoff multiplier e.g. value of 2 will double the delay, by default 1
+    logger : bool, optional
+        Option to log or print, by default True
+
+    Returns
+    -------
+    typing.Callable
+        Decorated callable that calls itself when exception(s) occur.
+
+    Examples
+    --------
+    >>> import random
+    >>> @retry(exception=Exception, n_tries=4)
+    ... def test_random(text):
+    ...    x = random.random()
+    ...    if x < 0.5:
+    ...        raise Exception("Fail")
+    ...    else:
+    ...        print("Success: ", text)
+    >>> test_random("It works!")
+    """
+
+    if func is None:
+        return partial(
+            retry,
+            exception=exception,
+            n_tries=n_tries,
+            delay=delay,
+            backoff=backoff,
+            logger=logger,
+        )
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        ntries, ndelay = n_tries, delay
+
+        while ntries > 1:
+            try:
+                return func(*args, **kwargs)
+            except exception as e:
+                msg = f"{str(e)}\nRetrying in {ndelay} seconds..."
+                if logger:
+                    log.warning(msg)
+                else:
+                    print(msg)
+                time.sleep(ndelay)
+
+                ntries -= 1
+                ndelay *= backoff
+
+        return func(*args, **kwargs)
+
+    return wrapper

influxspeedtest/config/__init__.py

@@ -7,4 +7,4 @@
 else:
     config = 'config.ini'
 
-config = ConfigManager(config)
+config = ConfigManager(config)

influxspeedtest/config/configmanager.py

@@ -3,7 +3,7 @@
 import sys
 
 
-class ConfigManager():
+class ConfigManager:
 
     def __init__(self, config):
         print('Loading Configuration File {}'.format(config))
@@ -32,6 +32,9 @@ def _load_config_values(self):
         self.influx_password = self.config['INFLUXDB'].get('Password', fallback='')
         self.influx_ssl = self.config['INFLUXDB'].getboolean('SSL', fallback=False)
         self.influx_verify_ssl = self.config['INFLUXDB'].getboolean('Verify_SSL', fallback=True)
+        self.influx_retry_delay = self.config['INFLUXDB'].getint('Retry_Delay', 5)
+        self.influx_retry_count = self.config['INFLUXDB'].getint('Retry_Count', 5)
+        self.influx_retry_backoff = self.config['INFLUXDB'].getint('Backoff', 1)
 
         # Logging
         self.logging_level = self.config['LOGGING'].get('Level', fallback='debug')