From 43805593e8bcfd64be86530ee39d4e42e6fdc7a2 Mon Sep 17 00:00:00 2001 From: avi robusta Date: Sun, 2 Jul 2023 20:21:32 +0300 Subject: added strategy metric override --- .../integrations/prometheus/metrics/base_metric.py | 30 ++++++++++++++++++++-- .../metrics_service/prometheus_metrics_service.py | 4 +-- 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()) -- cgit v1.2.3