From 1da622116d0fb8e1dfb4a8f0f7f53edfafb5367e Mon Sep 17 00:00:00 2001
From: Chico Venancio <francisco@argyle.com>
Date: Tue, 10 Dec 2024 09:04:25 -0300
Subject: [PATCH] feat: expose labels and annotations in the model (#369)

We would like to use labels and annotations to configure the strategy to
behave differently for some workloads. This will expose labels and
annotations in `object_data.annotations`/ `object_data.labels` to be
used in strategies.
---
 robusta_krr/core/integrations/kubernetes/__init__.py | 10 ++++++++++
 robusta_krr/core/models/objects.py                   |  4 ++++
 2 files changed, 14 insertions(+)

diff --git a/robusta_krr/core/integrations/kubernetes/__init__.py b/robusta_krr/core/integrations/kubernetes/__init__.py
index 78419fdc..1ede3409 100644
--- a/robusta_krr/core/integrations/kubernetes/__init__.py
+++ b/robusta_krr/core/integrations/kubernetes/__init__.py
@@ -201,6 +201,14 @@ def __build_scannable_object(
         namespace = item.metadata.namespace
         kind = kind or item.__class__.__name__[2:]
 
+        labels = {}
+        annotations = {}
+        if item.metadata.labels:
+            labels = item.metadata.labels
+
+        if item.metadata.annotations:
+            annotations = item.metadata.annotations
+
         obj = K8sObjectData(
             cluster=self.cluster,
             namespace=namespace,
@@ -209,6 +217,8 @@ def __build_scannable_object(
             container=container.name,
             allocations=ResourceAllocations.from_container(container),
             hpa=self.__hpa_list.get((namespace, kind, name)),
+            labels=labels,
+            annotations= annotations
         )
         obj._api_resource = item
         return obj
diff --git a/robusta_krr/core/models/objects.py b/robusta_krr/core/models/objects.py
index e4b400d9..b80cf041 100644
--- a/robusta_krr/core/models/objects.py
+++ b/robusta_krr/core/models/objects.py
@@ -46,6 +46,8 @@ class K8sObjectData(pd.BaseModel):
     kind: KindLiteral
     allocations: ResourceAllocations
     warnings: set[PodWarning] = set()
+    labels: Optional[dict[str, str]]
+    annotations: Optional[dict[str, str]]
 
     _api_resource = pd.PrivateAttr(None)
 
@@ -98,6 +100,8 @@ def split_into_batches(self, n: int) -> list[K8sObjectData]:
                 namespace=self.namespace,
                 kind=self.kind,
                 allocations=self.allocations,
+                labels=self.labels,
+                annotations=self.annotations,
             )
             for batch in batched(self.pods, n)
         ]