From c3dccb91fe5907a42a285b34443591d28cb344fc Mon Sep 17 00:00:00 2001
From: FrankFoerster24 <108520305+FrankFoerster24@users.noreply.github.com>
Date: Tue, 2 Apr 2024 18:35:33 +0200
Subject: Add Mimir support (#168)

* Add Mimir support

* Update url postfix for Mimir

* Fix Mimir-specific headers

---------

Co-authored-by: Pavel Zhukov <33721692+LeaveMyYard@users.noreply.github.com>
Co-authored-by: LeaveMyYard <zhukovpavel2001@gmail.com>
---
 robusta_krr/core/integrations/prometheus/loader.py |  4 +-
 .../metrics_service/mimir_metrics_service.py       | 47 ++++++++++++++++++++++
 .../metrics_service/prometheus_metrics_service.py  |  5 +++
 3 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 robusta_krr/core/integrations/prometheus/metrics_service/mimir_metrics_service.py

(limited to 'robusta_krr')

diff --git a/robusta_krr/core/integrations/prometheus/loader.py b/robusta_krr/core/integrations/prometheus/loader.py
index 82c7d74..db49392 100644
--- a/robusta_krr/core/integrations/prometheus/loader.py
+++ b/robusta_krr/core/integrations/prometheus/loader.py
@@ -16,13 +16,13 @@ from robusta_krr.core.models.objects import K8sObjectData, PodData
 from .metrics_service.prometheus_metrics_service import PrometheusMetricsService
 from .metrics_service.thanos_metrics_service import ThanosMetricsService
 from .metrics_service.victoria_metrics_service import VictoriaMetricsService
+from .metrics_service.mimir_metrics_service import MimirMetricsService
 
 if TYPE_CHECKING:
     from robusta_krr.core.abstract.strategies import BaseStrategy, MetricsPodData
 
 logger = logging.getLogger("krr")
 
-
 class PrometheusMetricsLoader:
     def __init__(self, *, cluster: Optional[str] = None) -> None:
         """
@@ -56,7 +56,7 @@ class PrometheusMetricsLoader:
             metrics_to_check = [PrometheusMetricsService]
         else:
             logger.info("No Prometheus URL is specified, trying to auto-detect a metrics service")
-            metrics_to_check = [VictoriaMetricsService, ThanosMetricsService, PrometheusMetricsService]
+            metrics_to_check = [VictoriaMetricsService, ThanosMetricsService, MimirMetricsService, PrometheusMetricsService]
 
         for metric_service_class in metrics_to_check:
             service_name = metric_service_class.name()
diff --git a/robusta_krr/core/integrations/prometheus/metrics_service/mimir_metrics_service.py b/robusta_krr/core/integrations/prometheus/metrics_service/mimir_metrics_service.py
new file mode 100644
index 0000000..ea3af57
--- /dev/null
+++ b/robusta_krr/core/integrations/prometheus/metrics_service/mimir_metrics_service.py
@@ -0,0 +1,47 @@
+from typing import Optional
+
+from kubernetes.client import ApiClient
+from prometrix import MetricsNotFound
+
+from robusta_krr.utils.service_discovery import MetricsServiceDiscovery
+
+from .prometheus_metrics_service import PrometheusMetricsService
+
+class MimirMetricsDiscovery(MetricsServiceDiscovery):
+    def find_metrics_url(self, *, api_client: Optional[ApiClient] = None) -> Optional[str]:
+        """
+        Finds the Mimir Metrics URL using selectors.
+        Args:
+            api_client (Optional[ApiClient]): A Kubernetes API client. Defaults to None.
+        Returns:
+            Optional[str]: The discovered Mimir Metrics URL, or None if not found.
+        """
+        return super().find_url(
+            selectors=[
+                "app.kubernetes.io/name=mimir,app.kubernetes.io/component=query-frontend",
+            ]
+        )
+
+
+class MimirMetricsService(PrometheusMetricsService):
+    """
+    A class for fetching metrics from Mimir Metrics.
+    """
+
+    service_discovery = MimirMetricsDiscovery
+    url_postfix = "/prometheus"
+    additional_headers = {"X-Scope-OrgID": "anonymous"}
+
+    def check_connection(self):
+        """
+        Checks the connection to Prometheus.
+        Raises:
+            MimirMetricsNotFound: If the connection to Mimir Metrics cannot be established.
+        """
+        try:
+            super().check_connection()
+        except MetricsNotFound as e:
+            # This is to clarify which metrics service had the issue and not say its a prometheus issue
+            raise MetricsNotFound(
+                f"Couldn't connect to Mimir Metrics found under {self.prometheus.url}\nCaused by {e.__class__.__name__}: {e})"
+            ) from e
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 ebcb835..423a29a 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
@@ -51,6 +51,8 @@ class PrometheusMetricsService(MetricsService):
     """
 
     service_discovery: type[MetricsServiceDiscovery] = PrometheusDiscovery
+    url_postfix: str = ""
+    additional_headers: dict[str, str] = {}
 
     def __init__(
         self,
@@ -86,9 +88,12 @@ class PrometheusMetricsService(MetricsService):
                 f"{self.name()} instance could not be found while scanning in {self.cluster} cluster."
             )
 
+        self.url += self.url_postfix
+
         logger.info(f"Using {self.name()} at {self.url} for cluster {cluster or 'default'}")
 
         headers = settings.prometheus_other_headers
+        headers |= self.additional_headers
 
         if self.auth_header:
             headers |= {"Authorization": self.auth_header}
-- 
cgit v1.2.3