From 77bda6bd5d1b49ca65610b04226052b147cc83d6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9F=D0=B0=D0=B2=D0=B5=D0=BB=20=D0=96=D1=83=D0=BA=D0=BE?=
 =?UTF-8?q?=D0=B2?= <33721692+LeaveMyYard@users.noreply.github.com>
Date: Wed, 24 May 2023 18:04:13 +0300
Subject: Improve old pods loading

---
 robusta_krr/core/integrations/prometheus/loader.py | 49 +++++++++++++---------
 robusta_krr/core/models/objects.py                 |  2 +-
 2 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/robusta_krr/core/integrations/prometheus/loader.py b/robusta_krr/core/integrations/prometheus/loader.py
index 1f0c14e..4946577 100644
--- a/robusta_krr/core/integrations/prometheus/loader.py
+++ b/robusta_krr/core/integrations/prometheus/loader.py
@@ -107,6 +107,9 @@ class PrometheusLoader(Configurable):
                 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 gather_data(
         self,
         object: K8sObjectData,
@@ -126,31 +129,39 @@ class PrometheusLoader(Configurable):
     async def add_historic_pods(self, object: K8sObjectData, period: datetime.timedelta) -> None:
         """Find pods that were already deleted, but still have some metrics in Prometheus"""
 
-        if len(object.pods) == 0:
-            return
-
         period_literal = f"{int(period.total_seconds()) // 60 // 24}d"
-        owner = await asyncio.to_thread(
-            self.prometheus.custom_query,
-            query=f'kube_pod_owner{{pod="{next(iter(object.pods)).name}"}}[{period_literal}]',
-        )
-
-        if owner == []:
-            return
-
-        owner = owner[0]["metric"]["owner_name"]
-
-        related_pods = await asyncio.to_thread(
-            self.prometheus.custom_query, query=f'kube_pod_owner{{owner_name="{owner}"}}[{period_literal}]'
+        pod_owners: list[str]
+        pod_owner_kind: str
+
+        if object.kind == "Deployment":
+            replicasets = await self.query(
+                "kube_replicaset_owner{"
+                f'owner_name="{object.name}", '
+                f'owner_kind="Deployment", '
+                f'namespace="{object.namespace}"'
+                "}"
+                f"[{period_literal}]"
+            )
+            pod_owners = [replicaset["metric"]["replicaset"] for replicaset in replicasets]
+            pod_owner_kind = "ReplicaSet"
+        else:
+            pod_owners = [object.name]
+            pod_owner_kind = object.kind
+
+        owners_regex = "|".join(pod_owners)
+        related_pods = await self.query(
+            "kube_pod_owner{"
+            f'owner_name=~"{owners_regex}", '
+            f'owner_kind="{pod_owner_kind}", '
+            f'namespace="{object.namespace}"'
+            "}"
+            f"[{period_literal}]"
         )
 
         current_pods = {p.name for p in object.pods}
 
         object.pods += [
-            PodData(
-                name=pod["metric"]["pod"],
-                deleted=True,
-            )
+            PodData(name=pod["metric"]["pod"], deleted=True)
             for pod in related_pods
             if pod["metric"]["pod"] not in current_pods
         ]
diff --git a/robusta_krr/core/models/objects.py b/robusta_krr/core/models/objects.py
index 21f5d61..3fd0dc8 100644
--- a/robusta_krr/core/models/objects.py
+++ b/robusta_krr/core/models/objects.py
@@ -19,7 +19,7 @@ class K8sObjectData(pd.BaseModel):
     container: str
     pods: list[PodData]
     namespace: str
-    kind: Optional[str]
+    kind: str
     allocations: ResourceAllocations
 
     def __str__(self) -> str:
-- 
cgit v1.2.3