Refactor HPA listing, add namespace to key

1 files changed, 37 insertions, 24 deletions

diff --git a/robusta_krr/core/integrations/kubernetes.py b/robusta_krr/core/integrations/kubernetes.py
index c947c80..1ffd3f4 100644
--- a/robusta_krr/core/integrations/kubernetes.py
+++ b/robusta_krr/core/integrations/kubernetes.py
@@ -29,6 +29,7 @@ from robusta_krr.utils.configurable import Configurable
 
 
 AnyKubernetesAPIObject = Union[V1Deployment, V1DaemonSet, V1StatefulSet, V1Pod, V1Job]
+HPAKey = tuple[str, str, str]
 
 
 class ClusterLoader(Configurable):
@@ -130,7 +131,7 @@ class ClusterLoader(Configurable):
             container=container.name,
             allocations=ResourceAllocations.from_container(container),
             pods=await self.__list_pods(item),
-            hpa=self.__hpa_list.get((kind, name)),
+            hpa=self.__hpa_list.get((namespace, kind, name)),
         )
 
     async def _list_deployments(self) -> list[K8sObjectData]:
@@ -211,7 +212,7 @@ class ClusterLoader(Configurable):
             *[self.__build_obj(item, container) for item in ret.items for container in item.spec.containers]
         )
 
-    async def __list_hpa_v1(self) -> dict[tuple[str, str], HPAData]:
+    async def __list_hpa_v1(self) -> dict[HPAKey, HPAData]:
         loop = asyncio.get_running_loop()
 
         res: V1HorizontalPodAutoscalerList = await loop.run_in_executor(
@@ -219,7 +220,11 @@ class ClusterLoader(Configurable):
         )
 
         return {
-            (hpa.spec.scale_target_ref.kind, hpa.spec.scale_target_ref.name): HPAData(
+            (
+                hpa.metadata.namespace,
+                hpa.spec.scale_target_ref.kind,
+                hpa.spec.scale_target_ref.name,
+            ): HPAData(
                 min_replicas=hpa.spec.min_replicas,
                 max_replicas=hpa.spec.max_replicas,
                 current_replicas=hpa.status.current_replicas,
@@ -230,28 +235,13 @@ class ClusterLoader(Configurable):
             for hpa in res.items
         }
 
-    async def __list_hpa(self) -> dict[tuple[str, str], HPAData]:
-        """List all HPA objects in the cluster.
-
-        Returns:
-            dict[tuple[str, str], HPAData]: A dictionary of HPA objects, indexed by scaleTargetRef (kind, name).
-        """
-
+    async def __list_hpa_v2(self) -> dict[HPAKey, HPAData]:
         loop = asyncio.get_running_loop()
 
-        try:
-            # Try to use V2 API first
-            res: V2HorizontalPodAutoscalerList = await loop.run_in_executor(
-                self.executor,
-                lambda: self.autoscaling_v2.list_horizontal_pod_autoscaler_for_all_namespaces(watch=False),
-            )
-        except ApiException as e:
-            if e.status != 404:
-                # If the error is other than not found, then re-raise it.
-                raise
-
-            # If V2 API does not exist, fall back to V1
-            return await self.__list_hpa_v1()
+        res: V2HorizontalPodAutoscalerList = await loop.run_in_executor(
+            self.executor,
+            lambda: self.autoscaling_v2.list_horizontal_pod_autoscaler_for_all_namespaces(watch=False),
+        )
 
         def __get_metric(hpa: V2HorizontalPodAutoscaler, metric_name: str) -> Optional[float]:
             return next(
@@ -264,7 +254,11 @@ class ClusterLoader(Configurable):
             )
 
         return {
-            (hpa.spec.scale_target_ref.kind, hpa.spec.scale_target_ref.name): HPAData(
+            (
+                hpa.metadata.namespace,
+                hpa.spec.scale_target_ref.kind,
+                hpa.spec.scale_target_ref.name,
+            ): HPAData(
                 min_replicas=hpa.spec.min_replicas,
                 max_replicas=hpa.spec.max_replicas,
                 current_replicas=hpa.status.current_replicas,
@@ -275,6 +269,25 @@ class ClusterLoader(Configurable):
             for hpa in res.items
         }
 
+    # TODO: What should we do in case of other metrics bound to the HPA?
+    async def __list_hpa(self) -> dict[HPAKey, HPAData]:
+        """List all HPA objects in the cluster.
+
+        Returns:
+            dict[tuple[str, str], HPAData]: A dictionary of HPA objects, indexed by scaleTargetRef (kind, name).
+        """
+
+        try:
+            # Try to use V2 API first
+            return await self.__list_hpa_v2()
+        except ApiException as e:
+            if e.status != 404:
+                # If the error is other than not found, then re-raise it.
+                raise
+
+            # If V2 API does not exist, fall back to V1
+            return await self.__list_hpa_v1()
+
 
 class KubernetesLoader(Configurable):
     async def list_clusters(self) -> Optional[list[str]]: