Bad cluster label exception

3 files changed, 28 insertions, 3 deletions

diff --git a/robusta_krr/core/integrations/prometheus/__init__.py b/robusta_krr/core/integrations/prometheus/__init__.py
index 3d298c6..aac0103 100644
--- a/robusta_krr/core/integrations/prometheus/__init__.py
+++ b/robusta_krr/core/integrations/prometheus/__init__.py
@@ -1,2 +1,2 @@
 from .loader import MetricsLoader
-from .metrics_service.prometheus_metrics_service import CustomPrometheusConnect, PrometheusDiscovery, PrometheusNotFound
+from .metrics_service.prometheus_metrics_service import CustomPrometheusConnect, PrometheusDiscovery, PrometheusNotFound, ClusterNotSpecifiedException
diff --git a/robusta_krr/core/integrations/prometheus/loader.py b/robusta_krr/core/integrations/prometheus/loader.py
index 773e8b7..cb33a4e 100644
--- a/robusta_krr/core/integrations/prometheus/loader.py
+++ b/robusta_krr/core/integrations/prometheus/loader.py
@@ -62,6 +62,7 @@ class MetricsLoader(Configurable):
                 loader = metric_service_class(config, api_client=api_client, cluster=cluster)
                 loader.check_connection()
                 self.echo(f"{service_name} found")
+                loader.validate_cluster_name()
                 return loader
             except MetricsNotFound as e:
                 self.debug(f"{service_name} not found")
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 7ceb766..ed1d65d 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
@@ -49,6 +49,12 @@ class PrometheusNotFound(MetricsNotFound):
 
     pass
 
+class ClusterNotSpecifiedException(Exception):
+    """
+    An exception raised when a prometheus requires a cluster label but an invalid one is provided.
+    """
+
+    pass
 
 class CustomPrometheusConnect(PrometheusConnect):
     """
@@ -131,11 +137,29 @@ class PrometheusMetricsService(MetricsService):
                 f"Couldn't connect to Prometheus found under {self.prometheus.url}\nCaused by {e.__class__.__name__}: {e})"
             ) from e
 
+    
     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")
+    def validate_cluster_name(self):
+        cluster_label = self.config.prometheus_cluster_label
+        cluster_names = self.get_cluster_names()
+    
+        if len(cluster_names) <= 1:
+            # there is only one cluster of metrics in this prometheus
+            return
+
+        if not cluster_label:
+            raise ClusterNotSpecifiedException(
+                f"KRR requires a prometheus_cluster_label to run your query, please select one of the following clusters {cluster_names}"
+                )
+        if cluster_label not in  cluster_names:
+            raise ClusterNotSpecifiedException(
+                f"prometheus_cluster_label {cluster_label} does not exist, please use one of the following {cluster_names}"
+                )
+
+    def get_cluster_names(self) -> Optional[List[str]]:
+        return self.prometheus.get_label_values(label_name="cluster")
 
     async def gather_data(
         self,