Merge pull request #45 from robusta-dev/pass-promql-to-result

Change format for metric, add start_time, end_time and step

5 files changed, 37 insertions, 27 deletions

diff --git a/robusta_krr/core/abstract/strategies.py b/robusta_krr/core/abstract/strategies.py
index 9c4f133..286c6c5 100644
--- a/robusta_krr/core/abstract/strategies.py
+++ b/robusta_krr/core/abstract/strategies.py
@@ -9,7 +9,7 @@ from textwrap import dedent
 
 import pydantic as pd
 
-from robusta_krr.core.models.result import K8sObjectData, ResourceType
+from robusta_krr.core.models.result import K8sObjectData, ResourceType, Metric
 from robusta_krr.utils.display_name import add_display_name
 
 
@@ -46,7 +46,7 @@ ArrayNx2 = Annotated[NDArray[np.float64], Literal["N", 2]]
 
 
 class ResourceHistoryData(pd.BaseModel):
-    query: str  # The query used to get the data
+    metric: Metric
     data: dict[str, ArrayNx2]  # Mapping: pod -> (time, value)
 
     class Config:
diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
index a409810..b0c7699 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_filtered_metric.py
@@ -1,7 +1,7 @@
-import datetime
 from typing import Any, Optional
 
 from .base_metric import BaseMetricLoader
+from robusta_krr.core.abstract.strategies import Metric
 
 PrometheusSeries = Any
 
@@ -45,7 +45,9 @@ class BaseFilteredMetricLoader(BaseMetricLoader):
         # takes kubelet job if exists, return first job alphabetically if it doesn't
         for target_name in target_names:
             relevant_series = [
-                series for series in series_list_result if BaseFilteredMetricLoader.get_target_name(series) == target_name
+                series
+                for series in series_list_result
+                if BaseFilteredMetricLoader.get_target_name(series) == target_name
             ]
             relevant_kubelet_metric = [series for series in relevant_series if series["metric"].get("job") == "kubelet"]
             if len(relevant_kubelet_metric) == 1:
@@ -55,8 +57,6 @@ class BaseFilteredMetricLoader(BaseMetricLoader):
             return_list.append(sorted_relevant_series[0])
         return return_list
 
-    async def query_prometheus(
-        self, query: str, start_time: datetime.datetime, end_time: datetime.datetime, step: datetime.timedelta
-    ) -> list[PrometheusSeries]:
-        result = await super().query_prometheus(query, start_time, end_time, step)
+    async def query_prometheus(self, metric: Metric) -> list[PrometheusSeries]:
+        result = await super().query_prometheus(metric)
         return self.filter_prom_jobs_results(result)
diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
index 2547d16..de4cf59 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
@@ -5,7 +5,7 @@ import asyncio
 import datetime
 from typing import TYPE_CHECKING, Callable, TypeVar
 
-from robusta_krr.core.abstract.strategies import ResourceHistoryData
+from robusta_krr.core.abstract.strategies import ResourceHistoryData, Metric
 from robusta_krr.core.models.config import Config
 from robusta_krr.core.models.objects import K8sObjectData
 from robusta_krr.utils.configurable import Configurable
@@ -27,37 +27,39 @@ class BaseMetricLoader(Configurable, abc.ABC):
     def get_query(self, object: K8sObjectData) -> str:
         ...
 
-    async def query_prometheus(
-        self, query: str, start_time: datetime.datetime, end_time: datetime.datetime, step: datetime.timedelta
-    ) -> list[dict]:
+    def _step_to_string(self, step: datetime.timedelta) -> str:
+        return f"{int(step.total_seconds()) // 60}m"
+
+    async def query_prometheus(self, metric: Metric) -> list[dict]:
         return await asyncio.to_thread(
             self.prometheus.custom_query_range,
-            query=query,
-            start_time=start_time,
-            end_time=end_time,
-            step=f"{int(step.total_seconds()) // 60}m",
+            query=metric.query,
+            start_time=metric.start_time,
+            end_time=metric.end_time,
+            step=metric.step,
         )
 
     async def load_data(
         self, object: K8sObjectData, period: datetime.timedelta, step: datetime.timedelta
     ) -> ResourceHistoryData:
         query = self.get_query(object)
-        result = await self.query_prometheus(
+        end_time = datetime.datetime.now()
+        metric = Metric(
             query=query,
-            start_time=datetime.datetime.now() - period,
-            end_time=datetime.datetime.now(),
-            step=step,
+            start_time=end_time - period,
+            end_time=end_time,
+            step=self._step_to_string(step),
         )
+        result = await self.query_prometheus(metric)
 
         if result == []:
             self.warning(f"Prometheus returned no {self.__class__.__name__} metrics for {object}")
-            return ResourceHistoryData(query=query, data={})
+            return ResourceHistoryData(metric=metric, data={})
 
         return ResourceHistoryData(
-            query=query,
+            metric=metric,
             data={
-                pod_result['metric']['pod']: np.array(pod_result["values"], dtype=np.float64)
-                for pod_result in result
+                pod_result["metric"]["pod"]: np.array(pod_result["values"], dtype=np.float64) for pod_result in result
             },
         )
 
diff --git a/robusta_krr/core/models/result.py b/robusta_krr/core/models/result.py
index 5119de8..0bad398 100644
--- a/robusta_krr/core/models/result.py
+++ b/robusta_krr/core/models/result.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import enum
 import itertools
 from typing import Any, Union, Optional
+from datetime import datetime
 
 import pydantic as pd
 
@@ -60,7 +61,14 @@ class ResourceRecommendation(pd.BaseModel):
     limits: dict[ResourceType, RecommendationValue]
 
 
-MetricsData = dict[ResourceType, str]
+class Metric(pd.BaseModel):
+    query: str
+    start_time: datetime
+    end_time: datetime
+    step: str
+
+
+MetricsData = dict[ResourceType, Metric]
 
 
 class ResourceScan(pd.BaseModel):
diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py
index 40a7b8e..eda3c5d 100644
--- a/robusta_krr/core/runner.py
+++ b/robusta_krr/core/runner.py
@@ -109,14 +109,14 @@ class Runner(Configurable):
             ]
         )
         data = dict(zip(ResourceType, data_tuple))
-        queries = {resource: data[resource].query for resource in ResourceType}
+        metrics = {resource: data[resource].metric for resource in ResourceType}
 
         self.__progressbar.progress()
 
         # NOTE: We run this in a threadpool as the strategy calculation might be CPU intensive
         # But keep in mind that numpy calcluations will not block the GIL
         result = await asyncio.to_thread(self._strategy.run, data, object)
-        return self._format_result(result), queries
+        return self._format_result(result), metrics
 
     async def _gather_objects_recommendations(self, objects: list[K8sObjectData]) -> list[tuple[ResourceAllocations, MetricsData]]:
         recommendations: list[tuple[RunResult, MetricsData]] = await asyncio.gather(