Added time range filter to prometheus queries

src/krkn_lib/prometheus/krkn_prometheus.py

@@ -1,6 +1,7 @@
 import logging
 import re
 import sys
+from datetime import datetime
 
 from prometheus_api_client import PrometheusConnect
 
@@ -33,27 +34,50 @@ def __init__(
             sys.exit(1)
 
     # Process custom prometheus query
-    def process_prom_query(self, query: str) -> list[dict[str:any]]:
+    def process_prom_query(
+        self,
+        query: str,
+        start_time: datetime = None,
+        end_time: datetime = None,
+        granularity: str = "1s",
+    ) -> list[dict[str:any]]:
         """
         Executes a query to the Prometheus API in PromQL language
 
         :param query: promQL query
+        :param start_time: start time of the result set (if None
+            no time filter is applied to the query)
+        :param end_time: end time of the result set (if None
+            no time filter is applied to the query)
+        :param granularity: if a time filter is applied represents the
+            frequency of the samples collected
+
         :return: a list of records in dictionary format
         """
         if self.prom_cli:
             try:
-                return self.prom_cli.custom_query(query=query, params=None)
+                if start_time is None or end_time is None:
+                    return self.prom_cli.custom_query(query=query, params=None)
+                else:
+                    return self.prom_cli.custom_query_range(
+                        query=query,
+                        start_time=start_time,
+                        end_time=end_time,
+                        step=granularity,
+                    )
             except Exception as e:
                 logging.error("Failed to get the metrics: %s" % e)
-                sys.exit(1)
+                raise e
         else:
             logging.info(
                 "Skipping the prometheus query as the "
                 "prometheus client couldn't "
                 "be initialized\n"
             )
 
-    def process_alert(self, alert: dict[str, str]):
+    def process_alert(
+        self, alert: dict[str, str], start_time: datetime, end_time: datetime
+    ):
         """
         Processes Krkn alarm in the format
 
@@ -87,6 +111,11 @@ def process_alert(self, alert: dict[str, str]):
 
         :params alert: a dictionary containing the following keys :
             expr, description, severity
+        :param start_time: start time of the result set (if None
+            no time filter is applied to the query)
+        :param end_time: end time of the result set (if None
+            no time filter is applied to the query)
+
 
 
         """
@@ -119,7 +148,9 @@ def process_alert(self, alert: dict[str, str]):
             return
 
         try:
-            records = self.process_prom_query(alert["expr"])
+            records = self.process_prom_query(
+                alert["expr"], start_time, end_time
+            )
             if len(records) == 0:
                 return