Skip to content

Commit

Permalink
Bad cluster label exception
Browse files Browse the repository at this point in the history
  • Loading branch information
Avi-Robusta committed Jun 13, 2023
1 parent d326e96 commit 16184a3
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 3 deletions.
2 changes: 1 addition & 1 deletion robusta_krr/core/integrations/prometheus/__init__.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions robusta_krr/core/integrations/prometheus/loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 16184a3

Please sign in to comment.