chore: use self.base_container_name to set self.container_logs value

apache · Jul 19, 2024 · c45c093 · c45c093
1 parent ada3b9a
commit c45c093
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py
@@ -30,7 +30,7 @@
 from contextlib import AbstractContextManager
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Iterable, Literal, Sequence
 
 import kubernetes
 import tenacity
@@ -91,7 +91,6 @@
 if TYPE_CHECKING:
     import jinja2
     from pendulum import DateTime
-    from typing_extensions import Literal
 
     from airflow.providers.cncf.kubernetes.secret import Secret
     from airflow.utils.context import Context
@@ -281,7 +280,8 @@ def __init__(
         startup_timeout_seconds: int = 120,
         startup_check_interval_seconds: int = 5,
         get_logs: bool = True,
-        container_logs: Iterable[str] | str | Literal[True] = BASE_CONTAINER_NAME,
+        base_container_name: str | None = None,
+        container_logs: Iterable[str] | str | Literal[True] | None = None,
         image_pull_policy: str | None = None,
         annotations: dict | None = None,
         container_resources: k8s.V1ResourceRequirements | None = None,
@@ -311,7 +311,6 @@ def __init__(
         termination_grace_period: int | None = None,
         configmaps: list[str] | None = None,
         skip_on_exit_code: int | Container[int] | None = None,
-        base_container_name: str | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         poll_interval: float = 2,
         log_pod_spec_on_failure: bool = True,
@@ -353,9 +352,9 @@ def __init__(
         self.cluster_context = cluster_context
         self.reattach_on_restart = reattach_on_restart
         self.get_logs = get_logs
-        self.container_logs = container_logs
-        if self.container_logs == KubernetesPodOperator.BASE_CONTAINER_NAME:
-            self.container_logs = base_container_name or self.BASE_CONTAINER_NAME
+        self.base_container_name = base_container_name or self.BASE_CONTAINER_NAME
+        # container_logs defaults to the provided base container name, or operator-level BASE_CONTAINER_NAME
+        self.container_logs = container_logs or self.base_container_name
         self.image_pull_policy = image_pull_policy
         self.node_selector = node_selector or {}
         self.annotations = annotations or {}
@@ -394,7 +393,6 @@ def __init__(
             if skip_on_exit_code is not None
             else []
         )
-        self.base_container_name = base_container_name or self.BASE_CONTAINER_NAME
         self.deferrable = deferrable
         self.poll_interval = poll_interval
         self.remote_pod: k8s.V1Pod | None = None