From 749ddb5e2355d3b38939237b9374823bd685b46d Mon Sep 17 00:00:00 2001
From: Charles Machalow <cmachalow@linkedin.com>
Date: Mon, 9 Jan 2023 09:20:52 -0800
Subject: [PATCH 1/4] Allow setting the name for the base container within K8s
 Pod Operator

Some downstream machinary may require a specific base container name
---
 .../kubernetes/operators/kubernetes_pod.py    |  15 ++-
 .../test_kubernetes_pod_operator.py           | 112 ++++++++++++++++++
 2 files changed, 121 insertions(+), 6 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
index c34fc5c02bbbf..e793700420b6c 100644
--- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
@@ -167,7 +167,7 @@ class KubernetesPodOperator(BaseOperator):
         during the next try. If False, always create a new pod for each try.
     :param labels: labels to apply to the Pod. (templated)
     :param startup_timeout_seconds: timeout in seconds to startup the pod.
-    :param get_logs: get the stdout of the container as logs of the tasks.
+    :param get_logs: get the stdout of the base container as logs of the tasks.
     :param image_pull_policy: Specify a policy to cache or always pull an image.
     :param annotations: non-identifying metadata you can attach to the Pod.
         Can be a large range of data, and can include characters
@@ -206,9 +206,10 @@ class KubernetesPodOperator(BaseOperator):
         to populate the environment variables with. The contents of the target
         ConfigMap's Data field will represent the key-value pairs as environment variables.
         Extends env_from.
+    :param base_container_name: The name of the base container in the pod. This container's logs
+        will appear as part of this task's logs if get_logs is True. Defaults to "base".
     """
 
-    BASE_CONTAINER_NAME = "base"
     POD_CHECKED_KEY = "already_checked"
 
     template_fields: Sequence[str] = (
@@ -272,6 +273,7 @@ def __init__(
         pod_runtime_info_envs: list[k8s.V1EnvVar] | None = None,
         termination_grace_period: int | None = None,
         configmaps: list[str] | None = None,
+        base_container_name: str = "base",
         **kwargs,
     ) -> None:
         # TODO: remove in provider 6.0.0 release. This is a mitigate step to advise users to switch to the
@@ -338,6 +340,7 @@ def __init__(
         self.termination_grace_period = termination_grace_period
         self.pod_request_obj: k8s.V1Pod | None = None
         self.pod: k8s.V1Pod | None = None
+        self.base_container_name = base_container_name
 
     @cached_property
     def _incluster_namespace(self):
@@ -498,12 +501,12 @@ def execute(self, context: Context):
             if self.get_logs:
                 self.pod_manager.fetch_container_logs(
                     pod=self.pod,
-                    container_name=self.BASE_CONTAINER_NAME,
+                    container_name=self.base_container_name,
                     follow=True,
                 )
             else:
                 self.pod_manager.await_container_completion(
-                    pod=self.pod, container_name=self.BASE_CONTAINER_NAME
+                    pod=self.pod, container_name=self.base_container_name
                 )
 
             if self.do_xcom_push:
@@ -535,7 +538,7 @@ def cleanup(self, pod: k8s.V1Pod, remote_pod: k8s.V1Pod):
             if self.log_events_on_failure:
                 self._read_pod_log_events(pod, reraise=False)
             self.process_pod_deletion(remote_pod, reraise=False)
-            error_message = get_container_termination_message(remote_pod, self.BASE_CONTAINER_NAME)
+            error_message = get_container_termination_message(remote_pod, self.base_container_name)
             raise AirflowException(
                 f"Pod {pod and pod.metadata.name} returned a failure:\n{error_message}\n"
                 f"remote_pod: {remote_pod}"
@@ -621,7 +624,7 @@ def build_pod_request_obj(self, context: Context | None = None) -> k8s.V1Pod:
                 containers=[
                     k8s.V1Container(
                         image=self.image,
-                        name=self.BASE_CONTAINER_NAME,
+                        name=self.base_container_name,
                         command=self.cmds,
                         ports=self.ports,
                         image_pull_policy=self.image_pull_policy,
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index caa0ad87b56b0..d21850e20c69a 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1153,3 +1153,115 @@ def test_using_resources(self):
                 do_xcom_push=False,
                 resources=resources,
             )
+
+    def test_changing_base_container_name_with_get_logs(self):
+        k = KubernetesPodOperator(
+            namespace="default",
+            image="ubuntu:16.04",
+            cmds=["bash", "-cx"],
+            arguments=["echo 10"],
+            labels=self.labels,
+            task_id=str(uuid4()),
+            in_cluster=False,
+            do_xcom_push=False,
+            get_logs=True,
+            base_container_name="apple-sauce",
+        )
+        assert k.base_container_name == "apple-sauce"
+        context = create_context(k)
+        with mock.patch.object(
+            k.pod_manager, "fetch_container_logs", wraps=k.pod_manager.fetch_container_logs
+        ) as mock_fetch_container_logs:
+            k.execute(context)
+
+        assert mock_fetch_container_logs.call_args[1]["container_name"] == "apple-sauce"
+        actual_pod = self.api_client.sanitize_for_serialization(k.pod)
+        self.expected_pod["spec"]["containers"][0]["name"] = "apple-sauce"
+        assert self.expected_pod["spec"] == actual_pod["spec"]
+
+    def test_changing_base_container_name_no_logs(self):
+        """
+        This test checks BOTH a modified base container name AND the get_logs=False flow..
+          and as a result, also checks that the flow works with fast containers
+          See #26796
+        """
+        k = KubernetesPodOperator(
+            namespace="default",
+            image="ubuntu:16.04",
+            cmds=["bash", "-cx"],
+            arguments=["echo 10"],
+            labels=self.labels,
+            task_id=str(uuid4()),
+            in_cluster=False,
+            do_xcom_push=False,
+            get_logs=False,
+            base_container_name="apple-sauce",
+        )
+        assert k.base_container_name == "apple-sauce"
+        context = create_context(k)
+        with mock.patch.object(
+            k.pod_manager, "await_container_completion", wraps=k.pod_manager.await_container_completion
+        ) as mock_await_container_completion:
+            k.execute(context)
+
+        assert mock_await_container_completion.call_args[1]["container_name"] == "apple-sauce"
+        actual_pod = self.api_client.sanitize_for_serialization(k.pod)
+        self.expected_pod["spec"]["containers"][0]["name"] = "apple-sauce"
+        assert self.expected_pod["spec"] == actual_pod["spec"]
+
+    def test_changing_base_container_name_no_logs_long(self):
+        """
+        Similar to test_changing_base_container_name_no_logs, but ensures that
+        pods running longer than 1 second work too. See #26796
+        """
+        k = KubernetesPodOperator(
+            namespace="default",
+            image="ubuntu:16.04",
+            cmds=["bash", "-cx"],
+            arguments=["sleep 3"],
+            labels=self.labels,
+            task_id=str(uuid4()),
+            in_cluster=False,
+            do_xcom_push=False,
+            get_logs=False,
+            base_container_name="apple-sauce",
+        )
+        assert k.base_container_name == "apple-sauce"
+        context = create_context(k)
+        with mock.patch.object(
+            k.pod_manager, "await_container_completion", wraps=k.pod_manager.await_container_completion
+        ) as mock_await_container_completion:
+            k.execute(context)
+
+        assert mock_await_container_completion.call_args[1]["container_name"] == "apple-sauce"
+        actual_pod = self.api_client.sanitize_for_serialization(k.pod)
+        self.expected_pod["spec"]["containers"][0]["name"] = "apple-sauce"
+        self.expected_pod["spec"]["containers"][0]["args"] = ["sleep 3"]
+        assert self.expected_pod["spec"] == actual_pod["spec"]
+
+    def test_changing_base_container_name_failure(self):
+        k = KubernetesPodOperator(
+            namespace="default",
+            image="ubuntu:16.04",
+            cmds=["exit"],
+            arguments=["1"],
+            labels=self.labels,
+            task_id=str(uuid4()),
+            in_cluster=False,
+            do_xcom_push=False,
+            base_container_name="apple-sauce",
+        )
+        assert k.base_container_name == "apple-sauce"
+        context = create_context(k)
+
+        class ShortCircuitException(Exception):
+            pass
+
+        with mock.patch(
+            "airflow.providers.cncf.kubernetes.operators.kubernetes_pod.get_container_termination_message",
+            side_effect=ShortCircuitException(),
+        ) as mock_get_container_termination_message:
+            with pytest.raises(ShortCircuitException):
+                k.execute(context)
+
+        assert mock_get_container_termination_message.call_args[0][1] == "apple-sauce"

From 3d8fcd81ff3a94568c5818d4a316a47754efb43e Mon Sep 17 00:00:00 2001
From: Charles Machalow <cmachalow@linkedin.com>
Date: Wed, 11 Jan 2023 14:41:57 -0800
Subject: [PATCH 2/4] Re-add support for BASE_CONTAINER_NAME as a class var

---
 .../kubernetes/operators/kubernetes_pod.py    | 11 ++++++++---
 .../test_kubernetes_pod_operator.py           | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
index e793700420b6c..4bae3a9abf020 100644
--- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
+++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
@@ -207,9 +207,14 @@ class KubernetesPodOperator(BaseOperator):
         ConfigMap's Data field will represent the key-value pairs as environment variables.
         Extends env_from.
     :param base_container_name: The name of the base container in the pod. This container's logs
-        will appear as part of this task's logs if get_logs is True. Defaults to "base".
+        will appear as part of this task's logs if get_logs is True. Defaults to None. If None,
+        will consult the class variable BASE_CONTAINER_NAME (which defaults to "base") for the base
+        container name to use.
     """
 
+    # This field can be overloaded at the instance level via base_container_name
+    BASE_CONTAINER_NAME = "base"
+
     POD_CHECKED_KEY = "already_checked"
 
     template_fields: Sequence[str] = (
@@ -273,7 +278,7 @@ def __init__(
         pod_runtime_info_envs: list[k8s.V1EnvVar] | None = None,
         termination_grace_period: int | None = None,
         configmaps: list[str] | None = None,
-        base_container_name: str = "base",
+        base_container_name: str | None = None,
         **kwargs,
     ) -> None:
         # TODO: remove in provider 6.0.0 release. This is a mitigate step to advise users to switch to the
@@ -340,7 +345,7 @@ def __init__(
         self.termination_grace_period = termination_grace_period
         self.pod_request_obj: k8s.V1Pod | None = None
         self.pod: k8s.V1Pod | None = None
-        self.base_container_name = base_container_name
+        self.base_container_name = base_container_name or self.BASE_CONTAINER_NAME
 
     @cached_property
     def _incluster_namespace(self):
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index d21850e20c69a..070efe6549135 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1265,3 +1265,22 @@ class ShortCircuitException(Exception):
                 k.execute(context)
 
         assert mock_get_container_termination_message.call_args[0][1] == "apple-sauce"
+
+    def test_base_container_name_init_precedence(self):
+        assert (
+            KubernetesPodOperator(base_container_name="apple-sauce", task_id=str(uuid4())).base_container_name
+            == "apple-sauce"
+        )
+        assert (
+            KubernetesPodOperator(task_id=str(uuid4())).base_container_name
+            == KubernetesPodOperator.BASE_CONTAINER_NAME
+        )
+
+        class MyK8SPodOperator(KubernetesPodOperator):
+            BASE_CONTAINER_NAME = "tomato-sauce"
+
+        assert (
+            MyK8SPodOperator(base_container_name="apple-sauce", task_id=str(uuid4())).base_container_name
+            == "apple-sauce"
+        )
+        assert MyK8SPodOperator(task_id=str(uuid4())).base_container_name == "tomato-sauce"

From 8ca8659205a96b76f6c96bab6be9766f5996799d Mon Sep 17 00:00:00 2001
From: Charles Machalow <csm10495@gmail.com>
Date: Mon, 23 Jan 2023 09:26:16 -0800
Subject: [PATCH 3/4] Update kubernetes_tests/test_kubernetes_pod_operator.py

Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
---
 kubernetes_tests/test_kubernetes_pod_operator.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index 070efe6549135..ae20d49e3ade6 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1181,9 +1181,9 @@ def test_changing_base_container_name_with_get_logs(self):
 
     def test_changing_base_container_name_no_logs(self):
         """
-        This test checks BOTH a modified base container name AND the get_logs=False flow..
-          and as a result, also checks that the flow works with fast containers
-          See #26796
+        This test checks BOTH a modified base container name AND the get_logs=False flow,
+        and as a result, also checks that the flow works with fast containers
+        See https://github.com/apache/airflow/issues/26796
         """
         k = KubernetesPodOperator(
             namespace="default",

From c6c3ad262a2691d63af4a272c79acb575b931343 Mon Sep 17 00:00:00 2001
From: Charles Machalow <csm10495@gmail.com>
Date: Mon, 23 Jan 2023 09:26:33 -0800
Subject: [PATCH 4/4] Update kubernetes_tests/test_kubernetes_pod_operator.py

Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com>
---
 kubernetes_tests/test_kubernetes_pod_operator.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index ae20d49e3ade6..104c319011bfe 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1212,7 +1212,8 @@ def test_changing_base_container_name_no_logs(self):
     def test_changing_base_container_name_no_logs_long(self):
         """
         Similar to test_changing_base_container_name_no_logs, but ensures that
-        pods running longer than 1 second work too. See #26796
+        pods running longer than 1 second work too.
+        See https://github.com/apache/airflow/issues/26796
         """
         k = KubernetesPodOperator(
             namespace="default",
