Add function to store metrics locally

Add function to print metric data
Add requirements.txt

Author: Anand Sanmukhani

prometheus_connect/__init__.py

@@ -1 +1,2 @@
+name = "prometheus_connect"
 from .prometheus_connect import PrometheusConnect

prometheus_connect/prometheus_connect.py

@@ -5,7 +5,9 @@
 import time
 import dateparser
 import sys
+import os
 from retrying import retry
+import bz2
 
 # set up logging
 import logging
@@ -55,7 +57,7 @@ def get_current_metric_value(self, metric_name, label_config = None):
             query = metric_name + "{" + ",".join(label_list) + "}"
         else:
             query = metric_name
-            
+
         response = requests.get('{0}/api/v1/query'.format(self.url),    # using the query API to get raw data
                                 params={'query': query},#label_config},
                                 verify=False, # Disable ssl certificate verification temporarily
@@ -72,7 +74,7 @@ def get_current_metric_value(self, metric_name, label_config = None):
         return (data)
 
     @retry(stop_max_attempt_number=MAX_REQUEST_RETRIES, wait_fixed=CONNECTION_RETRY_WAIT_TIME)
-    def get_metric_range_data(self, metric_name, start_time, end_time='now', chunk_size=None,label_config=None):
+    def get_metric_range_data(self, metric_name, start_time, end_time='now', chunk_size=None,label_config=None, store_locally=False):
         data = []
 
         start = int(dateparser.parse(str(start_time)).timestamp())
@@ -110,5 +112,44 @@ def get_metric_range_data(self, metric_name, start_time, end_time='now', chunk_s
                     requests.status_codes._codes[response.status_code][0],
                     response.content
                 ))
+            if store_locally:
+                # store it locally
+                self.store_metric_values_local(metric_name , (response.json()['data']['result']), start + chunk_seconds)
+
             start += chunk_seconds
         return (data)
+
+    def store_metric_values_local(self, metric_name, values, end_timestamp, file_path=None, compressed=True):
+        '''
+        Function to store metrics locally
+        '''
+        if not values:
+            return "No values for {}".format(metric_name)
+
+        if not file_path:
+            file_path = self._metric_filename(metric_name, end_timestamp)
+
+        if compressed:
+            payload = bz2.compress(str(values).encode('utf-8'))
+            file_path = file_path + ".bz2"
+        else:
+            payload = (str(values).encode('utf-8'))
+
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+        with open(file_path, "wb") as file:
+            file.write(payload)
+
+    def _metric_filename(self, metric_name, end_timestamp):
+        '''
+        Adds a timestamp to the filename before it is stored
+        '''
+        end_timestamp = dateparser.parse(str(end_timestamp))
+        directory_name = end_timestamp.strftime("%Y%m%d")
+        timestamp = end_timestamp.strftime("%Y%m%d%H%M")
+        object_path = "./metrics/" + self.prometheus_host + "/" + metric_name + "/" + directory_name + "/" + timestamp + ".json"
+        return object_path
+
+    def pretty_print_metric(self, metric_data):
+        data = metric_data
+        for metric in data:
+            print(json.dumps(metric, indent=4, sort_keys=True))

requirements.txt

@@ -0,0 +1,3 @@
+retrying
+requests
+dateparser