added cluster name param for centralized support

8 files changed, 44 insertions, 2 deletions

diff --git a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
index 7118d8c..b036218 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/base_metric.py
@@ -30,6 +30,17 @@ class BaseMetricLoader(Configurable, abc.ABC):
         super().__init__(config)
         self.prometheus = prometheus
 
+    def get_prometheus_cluster_label(self) -> str:
+        """
+        Generates the cluster label for querying a centralized Prometheus
+
+        Returns:
+        str: a promql safe label string for querying the cluster.
+        """
+        if self.config.prometheus_cluster_label is None:
+            return ''
+        return f', cluster="{self.config.prometheus_cluster_label}"'
+
     @abc.abstractmethod
     def get_query(self, object: K8sObjectData) -> str:
         """
diff --git a/robusta_krr/core/integrations/prometheus/metrics/cpu_metric.py b/robusta_krr/core/integrations/prometheus/metrics/cpu_metric.py
index 67ec98a..c00047a 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/cpu_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/cpu_metric.py
@@ -9,10 +9,12 @@ from .base_metric import bind_metric
 class CPUMetricLoader(BaseFilteredMetricLoader):
     def get_query(self, object: K8sObjectData) -> str:
         pods_selector = "|".join(pod.name for pod in object.pods)
+        cluster_label = self.get_prometheus_cluster_label()
         return (
             "sum(irate(container_cpu_usage_seconds_total{"
             f'namespace="{object.namespace}", '
             f'pod=~"{pods_selector}", '
             f'container="{object.container}"'
+            f'{cluster_label}'
             "}[5m])) by (container, pod, job)"
         )
diff --git a/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py b/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
index ddb056a..dfbbc08 100644
--- a/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
+++ b/robusta_krr/core/integrations/prometheus/metrics/memory_metric.py
@@ -9,10 +9,12 @@ from .base_metric import bind_metric
 class MemoryMetricLoader(BaseFilteredMetricLoader):
     def get_query(self, object: K8sObjectData) -> 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)"
         )
diff --git a/robusta_krr/core/integrations/prometheus/metrics_service/base_metric_service.py b/robusta_krr/core/integrations/prometheus/metrics_service/base_metric_service.py
index 21b3706..0b41f84 100644
--- a/robusta_krr/core/integrations/prometheus/metrics_service/base_metric_service.py
+++ b/robusta_krr/core/integrations/prometheus/metrics_service/base_metric_service.py
@@ -1,6 +1,6 @@
 import abc
 import datetime
-from typing import Optional
+from typing import Optional, List
 
 from kubernetes.client.api_client import ApiClient
 
@@ -39,6 +39,10 @@ class MetricsService(Configurable, abc.ABC):
         return classname.replace("MetricsService", "") if classname != MetricsService.__name__ else classname
 
     @abc.abstractmethod
+    async def get_cluster_names(self) -> Optional[List[str]]:
+        ...
+
+    @abc.abstractmethod
     async def gather_data(
         self,
         object: K8sObjectData,
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 c5d42ef..7ceb766 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
@@ -1,6 +1,6 @@
 import asyncio
 import datetime
-from typing import Optional, no_type_check
+from typing import Optional, no_type_check, List
 
 import requests
 from kubernetes.client import ApiClient
@@ -134,6 +134,9 @@ class PrometheusMetricsService(MetricsService):
     async def query(self, query: str) -> dict:
         return await asyncio.to_thread(self.prometheus.custom_query, query=query)
 
+    async def get_cluster_names(self) -> Optional[List[str]]:
+        return await asyncio.to_thread(self.prometheus.get_label_values, label_name="cluster")
+
     async def gather_data(
         self,
         object: K8sObjectData,
diff --git a/robusta_krr/core/models/config.py b/robusta_krr/core/models/config.py
index 6cd0b56..72cf4d8 100644
--- a/robusta_krr/core/models/config.py
+++ b/robusta_krr/core/models/config.py
@@ -25,6 +25,7 @@ class Config(pd.BaseSettings):
     prometheus_url: Optional[str] = pd.Field(None)
     prometheus_auth_header: Optional[str] = pd.Field(None)
     prometheus_ssl_enabled: bool = pd.Field(False)
+    prometheus_cluster_label: Optional[str] = pd.Field(None)
 
     # Logging Settings
     format: str
diff --git a/robusta_krr/core/runner.py b/robusta_krr/core/runner.py
index 7a64bfb..28fd898 100644
--- a/robusta_krr/core/runner.py
+++ b/robusta_krr/core/runner.py
@@ -13,6 +13,12 @@ from robusta_krr.utils.logo import ASCII_LOGO
 from robusta_krr.utils.progress_bar import ProgressBar
 from robusta_krr.utils.version import get_version
 
+class ClusterNotSpecifiedException(Exception):
+    """
+    An exception raised when a cluster is not specified.
+    """
+    pass
+
 
 class Runner(Configurable):
     EXPECTED_EXCEPTIONS = (KeyboardInterrupt, PrometheusNotFound)
@@ -139,6 +145,11 @@ class Runner(Configurable):
 
     async def _collect_result(self) -> Result:
         clusters = await self._k8s_loader.list_clusters()
+        if len(clusters) > 1 and self.config.prometheus_url:
+            # this can only happen for multi-cluster querying a single centeralized prometheus
+            # In this scenario we dont yet support determining which metrics belong to which cluster so the reccomendation can be incorrect
+            raise ClusterNotSpecifiedException(f"Must specify only one cluster from {clusters} to for prometheus at {self.config.prometheus_url}")
+        
         self.info(f'Using clusters: {clusters if clusters is not None else "inner cluster"}')
         objects = await self._k8s_loader.list_scannable_objects(clusters)
 
diff --git a/robusta_krr/main.py b/robusta_krr/main.py
index dbf12d8..33480e9 100644
--- a/robusta_krr/main.py
+++ b/robusta_krr/main.py
@@ -91,6 +91,13 @@ def load_commands() -> None:
                     help="Enable SSL for Prometheus requests.",
                     rich_help_panel="Prometheus Settings",
                 ),
+                prometheus_cluster_label: Optional[str] = typer.Option(
+                    None,
+                    "--prometheus-cluster-label",
+                    "-l",
+                    help="The label in prometheus for your cluster.(Only relevant for centralized prometheus)",
+                    rich_help_panel="Prometheus Settings",
+                ),
                 format: str = typer.Option("table", "--formatter", "-f", help="Output formatter ({formatters})", rich_help_panel="Logging Settings"),
                 verbose: bool = typer.Option(False, "--verbose", "-v", help="Enable verbose mode", rich_help_panel="Logging Settings"),
                 quiet: bool = typer.Option(False, "--quiet", "-q", help="Enable quiet mode", rich_help_panel="Logging Settings"),
@@ -106,6 +113,7 @@ def load_commands() -> None:
                     prometheus_url=prometheus_url,
                     prometheus_auth_header=prometheus_auth_header,
                     prometheus_ssl_enabled=prometheus_ssl_enabled,
+                    prometheus_cluster_label=prometheus_cluster_label,
                     format=format,
                     verbose=verbose,
                     quiet=quiet,