Implement saving of total updated counter

closes #7
gmt2001 · Mar 7, 2022 · 178453f · 178453f
1 parent 741e8df
commit 178453f
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 3 deletions.
diff --git a/pyouroboros/config.py b/pyouroboros/config.py
@@ -9,7 +9,7 @@ class Config(object):
                'INFLUX_URL', 'INFLUX_PORT', 'INFLUX_USERNAME', 'INFLUX_PASSWORD', 'INFLUX_DATABASE', 'INFLUX_SSL',
                'INFLUX_VERIFY_SSL', 'DATA_EXPORT', 'SELF_UPDATE', 'LABEL_ENABLE', 'DOCKER_TLS', 'LABELS_ONLY',
                'DRY_RUN', 'MONITOR_ONLY', 'HOSTNAME', 'DOCKER_TLS_VERIFY', 'SWARM', 'SKIP_STARTUP_NOTIFICATIONS', 'LANGUAGE',
-               'TZ', 'CLEANUP_UNUSED_VOLUMES', 'DOCKER_TIMEOUT', 'LATEST_ONLY']
+               'TZ', 'CLEANUP_UNUSED_VOLUMES', 'DOCKER_TIMEOUT', 'LATEST_ONLY', 'SAVE_COUNTERS']
 
     hostname = environ.get('HOSTNAME')
     interval = 300
@@ -35,6 +35,8 @@ class Config(object):
     language = 'en'
     tz = 'UTC'
 
+    save_counters = False
+
     repo_user = None
     repo_pass = None
     auth_json = None

diff --git a/pyouroboros/dataexporters.py b/pyouroboros/dataexporters.py
@@ -1,9 +1,11 @@
 import prometheus_client
-
+import json
+from os import unlink
 from logging import getLogger
 from influxdb import InfluxDBClient
 from datetime import datetime, timezone
-
+from pathlib import Path
+from pyouroboros.helpers import get_exec_dir
 
 class DataManager(object):
     def __init__(self, config):
@@ -31,6 +33,25 @@ def set(self, socket):
         if self.config.data_export == "prometheus" and self.enabled:
             self.prometheus.set_monitored(socket)
 
+    def save(self):
+        if self.config.save_counters:
+            fpath = Path(get_exec_dir() + '/hooks/datamanager.json')
+            try:
+                with open(fpath, 'w') as file:
+                    json.dump(self.total_updated, file)
+            except:
+                self.logger.debug('Unable to save JSON')
+
+    def load(self):
+        if self.config.save_counters:
+            fpath = Path(get_exec_dir() + '/hooks/datamanager.json')
+            try:
+                with open(fpath, 'r') as file:
+                    self.total_updated = json.load(file)
+                unlink(fpath)
+            except:
+                self.logger.debug('No JSON to load or unlink failed')
+
 
 class PrometheusExporter(object):
     def __init__(self, data_manager, config):

diff --git a/pyouroboros/dockerclient.py b/pyouroboros/dockerclient.py
@@ -433,11 +433,13 @@ def update_self(self, count=None, old_container=None, me_list=None, new_image=No
             self.logger.debug('Ahhh. All better.')
             run_hook('after_self_cleanup', None, locals)
 
+            self.data_manager.load()
             self.monitored = self.monitor_filter()
         elif count == 1:
             self.logger.debug('I need to update! Starting the ouroboros ;)')
             self_name = 'ouroboros-updated' if old_container.name == 'ouroboros' else 'ouroboros'
             new_config = set_properties(old=old_container, new=new_image, self_name=self_name)
+            self.data_manager.save()
             locals = {}
             locals['self_name'] = self_name
             locals['old_container'] = old_container

diff --git a/pyouroboros/ouroboros.py b/pyouroboros/ouroboros.py
@@ -114,6 +114,9 @@ def main():
                                    'EXAMPLE: MyPa$$w0rd')
 
     data_group = parser.add_argument_group('Data Export', 'Configuration of data export functionality')
+    data_group.add_argument('-sc', '--save-counters', default=Config.save_counters, dest='SAVE_COUNTERS',
+                            action='store_true', help='Save total-updated counters across self-updates')
+
     data_group.add_argument('-D', '--data-export', choices=['prometheus', 'influxdb'], default=Config.data_export,
                             dest='DATA_EXPORT', help='Enable exporting of data for chosen option')