PrometeusClusterLoader.list_clusters implementation

1 files changed, 16 insertions, 5 deletions

diff --git a/robusta_krr/core/integrations/prometheus/cluster_loader/__init__.py b/robusta_krr/core/integrations/prometheus/cluster_loader/__init__.py
index c073ab5..1fbe470 100644
--- a/robusta_krr/core/integrations/prometheus/cluster_loader/__init__.py
+++ b/robusta_krr/core/integrations/prometheus/cluster_loader/__init__.py
@@ -36,12 +36,22 @@ class PrometheusClusterLoader(BaseClusterLoader):
         self._prometheus_connector.connect(settings.prometheus_url)
 
     async def list_clusters(self) -> Optional[list[str]]:
-        if settings.prometheus_cluster_label is None:
+        if settings.prometheus_label is None:
+            logger.info(
+                "Assuming that Prometheus contains only one cluster."
+            )
+            logger.info("If you have multiple clusters in Prometheus, please provide the `-l` flag.")
             return None
 
-        # TODO: We can try to auto-discover clusters by querying Prometheus,
-        # but for that we will need to rework PrometheusMetric.get_prometheus_cluster_label
-        return [settings.prometheus_cluster_label]
+        clusters = await self.prometheus.loader.query(
+            f"""
+                avg by({settings.prometheus_label}) (
+                    kube_pod_container_resource_limits
+                )
+            """
+        )
+
+        return [cluster["metric"][settings.prometheus_label] for cluster in clusters["data"]["result"]]
 
     @cache
     def get_workload_loader(self, cluster: str) -> PrometheusWorkloadLoader:
@@ -75,4 +85,5 @@ class PrometheusWorkloadLoader(BaseWorkloadLoader):
 
         return workloads
 
-__all__ = ["PrometheusClusterLoader", "PrometheusWorkloadLoader"]
\ No newline at end of file
+
+__all__ = ["PrometheusClusterLoader", "PrometheusWorkloadLoader"]