added strategy metric override

2 files changed, 32 insertions, 10 deletions

diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
index 288a30f..e62f9ac 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
@@ -14,15 +14,16 @@ from robusta_krr.utils.configurable import Configurable
 
 if TYPE_CHECKING:
     from .. import CustomPrometheusConnect
+    MetricsDictionary = dict[str, type[BaseMetricLoader]]
 
 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: MetricsDictionary = {}
-STRATEGY_METRICS_OVERRIDES: dict[str,MetricsDictionary]
+STRATEGY_METRICS_OVERRIDES: dict[str,MetricsDictionary] = {}
 
 
 
@@ -172,8 +173,9 @@ 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]
+            lower_strategy = strategy.lower()
+            if lower_strategy and lower_strategy in STRATEGY_METRICS_OVERRIDES and resource in STRATEGY_METRICS_OVERRIDES[lower_strategy]:
+                return STRATEGY_METRICS_OVERRIDES[lower_strategy][resource]
             return REGISTERED_METRICS[resource]
         except KeyError as e:
             raise KeyError(f"Resource {resource} was not registered by `@bind_metric(...)`") from e
@@ -212,9 +214,10 @@ def override_metric(strategy: str, resource: str) -> Callable[[type[Self]], type
     """
 
     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
+        lower_strategy = strategy.lower()
+        if lower_strategy not in STRATEGY_METRICS_OVERRIDES:
+            STRATEGY_METRICS_OVERRIDES[lower_strategy] = {}
+        STRATEGY_METRICS_OVERRIDES[lower_strategy][resource] = cls
         return cls
 
     return decorator
\ No newline at end of file
diff --git a/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py b/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
index 2dcbed5..8f4e286 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
@@ -3,14 +3,33 @@ from robusta_krr.core.models.allocations import ResourceType
 from robusta_krr.core.models.objects import K8sObjectData
 
 from .base_filtered_metric import BaseFilteredMetricLoader
-from .base_metric import bind_metric, QueryType
-
+from .base_metric import bind_metric, QueryType, override_metric
 
 @bind_metric(ResourceType.Memory)
 class MemoryMetricLoader(BaseFilteredMetricLoader):
     def get_query(self, object: K8sObjectData, resolution: Optional[str]) -> str:
         pods_selector = "|".join(pod.name for pod in object.pods)
         cluster_label = self.get_prometheus_cluster_label()
+        return (
+            "sum(container_memory_working_set_bytes{"
+            f'namespace="{object.namespace}", '
+            f'pod=~"{pods_selector}", '
+            f'container="{object.container}"'
+            f"{cluster_label}"
+            "}) by (container, pod, job)"
+        )
+
+    def get_query_type(self) -> QueryType:
+        return QueryType.QueryRange
+
+@override_metric("simple", ResourceType.Memory)
+class MemoryMetricLoader(BaseFilteredMetricLoader):
+    """
+    A class that overrides the memory metric on the simple strategy.
+    """
+    def get_query(self, object: K8sObjectData, resolution: Optional[str]) -> str:
+        pods_selector = "|".join(pod.name for pod in object.pods)
+        cluster_label = self.get_prometheus_cluster_label()
         resolution_formatted = f'[{resolution}]' if resolution else ''
         return (
             f'max(max_over_time(container_memory_working_set_bytes{{'
@@ -23,4 +42,4 @@ class MemoryMetricLoader(BaseFilteredMetricLoader):
         )
     
     def get_query_type(self) -> QueryType:
-        return QueryType.Query
\ No newline at end of file
+        return QueryType.Query