diff --git a/robusta_krr/core/integrations/prometheus/__init__.py b/robusta_krr/core/integrations/prometheus/__init__.py index 3d298c6e..aac0103b 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 773e8b79..cb33a4ec 100644 --- a/robusta_krr/core/integrations/prometheus/loader.py +++ b/robusta_krr/core/integrations/prometheus/loader.py @@ -62,6 +62,7 @@ def get_metrics_service( 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 7ceb7664..ed1d65d8 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 @@ def check_connection(self): 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,