You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ka...@apache.org on 2023/02/01 11:51:09 UTC

[airflow] branch main updated: Patch only single label when marking KPO checked (#29279)

This is an automated email from the ASF dual-hosted git repository.

kaxilnaik pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new d26dc22391 Patch only single label when marking KPO checked (#29279)
d26dc22391 is described below

commit d26dc223915c50ff58252a709bb7b33f5417dfce
Author: Jed Cunningham <66...@users.noreply.github.com>
AuthorDate: Wed Feb 1 05:50:58 2023 -0600

    Patch only single label when marking KPO checked (#29279)
    
    Instead of sending over a whole pod, we can simply send over the label
    to add to the pod instead. This is less work for us, and also avoids any
    issues like #24015.
    
    Fixes #24015
---
 .../providers/cncf/kubernetes/operators/kubernetes_pod.py    |  8 +++++---
 .../cncf/kubernetes/operators/test_kubernetes_pod.py         | 12 ++++++++++++
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
index 2768c5ff1a..1874bf90f5 100644
--- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
@@ -711,9 +711,11 @@ class KubernetesPodOperator(BaseOperator):
     def patch_already_checked(self, pod: k8s.V1Pod, *, reraise=True):
         """Add an "already checked" annotation to ensure we don't reattach on retries"""
         with _optionally_suppress(reraise=reraise):
-            pod.metadata.labels[self.POD_CHECKED_KEY] = "True"
-            body = PodGenerator.serialize_pod(pod)
-            self.client.patch_namespaced_pod(pod.metadata.name, pod.metadata.namespace, body)
+            self.client.patch_namespaced_pod(
+                name=pod.metadata.name,
+                namespace=pod.metadata.namespace,
+                body={"metadata": {"labels": {self.POD_CHECKED_KEY: "True"}}},
+            )
 
     def on_kill(self) -> None:
         if self.pod:
diff --git a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py
index b04947a765..261d62dbe0 100644
--- a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py
+++ b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py
@@ -999,6 +999,18 @@ class TestKubernetesPodOperator:
         mock_patch_already_checked.assert_called_once()
         mock_delete_pod.assert_not_called()
 
+    @patch(HOOK_CLASS, new=MagicMock)
+    def test_patch_already_checked(self):
+        """Make sure we patch the pods with the right label"""
+        k = KubernetesPodOperator(task_id="task")
+        pod = k.build_pod_request_obj()
+        k.patch_already_checked(pod)
+        k.client.patch_namespaced_pod.assert_called_once_with(
+            name=pod.metadata.name,
+            namespace=pod.metadata.namespace,
+            body={"metadata": {"labels": {"already_checked": "True"}}},
+        )
+
     def test_task_id_as_name(self):
         k = KubernetesPodOperator(
             task_id=".hi.-_09HI",