added strategy metric override

2 files changed, 29 insertions, 5 deletions

diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
index c0acf89..288a30f 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
@@ -19,8 +19,11 @@ class QueryType(str, enum.Enum):
     Query="query"
     QueryRange="query_range"
 
+MetricsDictionary = dict[str, type[BaseMetricLoader]]
 # A registry of metrics that can be used to fetch a corresponding metric loader.
-REGISTERED_METRICS: dict[str, type[BaseMetricLoader]] = {}
+REGISTERED_METRICS: MetricsDictionary = {}
+STRATEGY_METRICS_OVERRIDES: dict[str,MetricsDictionary]
+
 
 
 class BaseMetricLoader(Configurable, abc.ABC):
@@ -153,12 +156,13 @@ class BaseMetricLoader(Configurable, abc.ABC):
         )
 
     @staticmethod
-    def get_by_resource(resource: str) -> type[BaseMetricLoader]:
+    def get_by_resource(resource: str, strategy: Optional[str]) -> type[BaseMetricLoader]:
         """
         Fetches the metric loader corresponding to the specified resource.
 
         Args:
         resource (str): The name of the resource.
+        resource (str): The name of the strategy.
 
         Returns:
         type[BaseMetricLoader]: The class of the metric loader corresponding to the resource.
@@ -168,6 +172,8 @@ class BaseMetricLoader(Configurable, abc.ABC):
         """
 
         try:
+            if strategy and STRATEGY_METRICS_OVERRIDES.has_key(strategy) and STRATEGY_METRICS_OVERRIDES[strategy].has_key(resource):
+                return STRATEGY_METRICS_OVERRIDES[strategy][resource]
             return REGISTERED_METRICS[resource]
         except KeyError as e:
             raise KeyError(f"Resource {resource} was not registered by `@bind_metric(...)`") from e
@@ -192,3 +198,23 @@ def bind_metric(resource: str) -> Callable[[type[Self]], type[Self]]:
         return cls
 
     return decorator
+
+def override_metric(strategy: str, resource: str) -> Callable[[type[Self]], type[Self]]:
+    """
+    A decorator that overrides the bound metric on a specific strategy.
+
+    Args:
+    strategy (str): The name of the strategy for this metric.
+    resource (str): The name of the resource.
+
+    Returns:
+    Callable[[type[Self]], type[Self]]: The decorator that does the binding.
+    """
+
+    def decorator(cls: type[Self]) -> type[Self]:
+        if not STRATEGY_METRICS_OVERRIDES.has_key(strategy):
+            STRATEGY_METRICS_OVERRIDES[strategy] = {}
+        STRATEGY_METRICS_OVERRIDES[strategy][resource] = cls
+        return cls
+
+    return decorator
\ No newline at end of file
diff --git a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
index 896387f..919b508 100644
--- a/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
+++ b/robusta_krr/core/integrations/prometheus/metrics_service/prometheus_metrics_service.py
@@ -177,9 +177,7 @@ class PrometheusMetricsService(MetricsService):
         """
         self.debug(f"Gathering data for {object} and {resource}")
 
-        await self.add_historic_pods(object, period)
-
-        MetricLoaderType = BaseMetricLoader.get_by_resource(resource)
+        MetricLoaderType = BaseMetricLoader.get_by_resource(resource, self.config.strategy)
         metric_loader = MetricLoaderType(self.config, self.prometheus, self.executor)
         return await metric_loader.load_data(object, period, step, self.name())