diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py index 564cd9f394156..c2d9e540a505a 100644 --- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py +++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py @@ -73,8 +73,11 @@ class KubernetesPodOperator(BaseOperator): but fully qualified URLS will point to custom repositories. (templated) :type image: str :param name: name of the pod in which the task will run, will be used (plus a random - suffix) to generate a pod id (DNS-1123 subdomain, containing only [a-z0-9.-]). + suffix if random_name_suffix is True) to generate a pod id (DNS-1123 subdomain, + containing only [a-z0-9.-]). :type name: str + :param random_name_suffix: if True, will generate a random suffix. + :type random_name_suffix: bool :param cmds: entrypoint of the container. (templated) The docker images's entrypoint is used if this is not provided. :type cmds: list[str] @@ -180,6 +183,7 @@ def __init__( namespace: Optional[str] = None, image: Optional[str] = None, name: Optional[str] = None, + random_name_suffix: Optional[bool] = True, cmds: Optional[List[str]] = None, arguments: Optional[List[str]] = None, ports: Optional[List[k8s.V1ContainerPort]] = None, @@ -275,6 +279,7 @@ def __init__( self.priority_class_name = priority_class_name self.pod_template_file = pod_template_file self.name = self._set_name(name) + self.random_name_suffix = random_name_suffix self.termination_grace_period = termination_grace_period self.client: CoreV1Api = None self.pod: k8s.V1Pod = None @@ -450,7 +455,7 @@ def create_pod_request_obj(self) -> k8s.V1Pod: metadata=k8s.V1ObjectMeta( namespace=self.namespace, labels=self.labels, - name=PodGenerator.make_unique_pod_id(self.name), + name=self.name, annotations=self.annotations, ), spec=k8s.V1PodSpec( @@ -486,6 +491,9 @@ def create_pod_request_obj(self) -> k8s.V1Pod: pod = PodGenerator.reconcile_pods(pod_template, pod) + if self.random_name_suffix: + pod.metadata.name = PodGenerator.make_unique_pod_id(pod.metadata.name) + for secret in self.secrets: self.log.debug("Adding secret to task %s", self.task_id) pod = secret.attach_to_pod(pod) diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py index 52d7f3ab7bec7..d51a5e0633d01 100644 --- a/kubernetes_tests/test_kubernetes_pod_operator.py +++ b/kubernetes_tests/test_kubernetes_pod_operator.py @@ -824,7 +824,10 @@ def test_pod_template_file(self, mock_client, monitor_mock, start_mock): path = sys.path[0] + '/tests/kubernetes/pod.yaml' k = KubernetesPodOperator( - task_id="task" + self.get_current_task_name(), pod_template_file=path, do_xcom_push=True + task_id="task" + self.get_current_task_name(), + random_name_suffix=False, + pod_template_file=path, + do_xcom_push=True, ) monitor_mock.return_value = (State.SUCCESS, None, None) diff --git a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py index 7f4e53bc19b3a..dfb99ed3e0850 100644 --- a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py @@ -20,6 +20,7 @@ import pytest from kubernetes.client import ApiClient, models as k8s +from parameterized import parameterized from airflow.exceptions import AirflowException from airflow.models import DAG, DagRun, TaskInstance @@ -209,13 +210,15 @@ def test_pod_delete_even_on_launcher_error(self, delete_pod_mock): k.execute(context=context) assert delete_pod_mock.called - def test_randomize_pod_name(self): + @parameterized.expand([[True], [False]]) + def test_provided_pod_name(self, randomize_name): name_base = "test" k = KubernetesPodOperator( namespace="default", image="ubuntu:16.04", name=name_base, + random_name_suffix=randomize_name, task_id="task", in_cluster=False, do_xcom_push=False, @@ -223,8 +226,11 @@ def test_randomize_pod_name(self): ) pod = k.create_pod_request_obj() - assert pod.metadata.name.startswith(name_base) - assert pod.metadata.name != name_base + if randomize_name: + assert pod.metadata.name.startswith(name_base) + assert pod.metadata.name != name_base + else: + assert pod.metadata.name == name_base def test_pod_name_required(self): with pytest.raises(AirflowException, match="`name` is required"): @@ -237,9 +243,13 @@ def test_pod_name_required(self): cluster_context="default", ) - def test_full_pod_spec(self): + @parameterized.expand([[True], [False]]) + def test_full_pod_spec(self, randomize_name): + pod_spec_name_base = "hello" pod_spec = k8s.V1Pod( - metadata=k8s.V1ObjectMeta(name="hello", labels={"foo": "bar"}, namespace="mynamespace"), + metadata=k8s.V1ObjectMeta( + name=pod_spec_name_base, labels={"foo": "bar"}, namespace="mynamespace" + ), spec=k8s.V1PodSpec( containers=[ k8s.V1Container( @@ -253,6 +263,7 @@ def test_full_pod_spec(self): k = KubernetesPodOperator( task_id="task", + random_name_suffix=randomize_name, in_cluster=False, do_xcom_push=False, cluster_context="default", @@ -260,7 +271,11 @@ def test_full_pod_spec(self): ) pod = self.run_pod(k) - assert pod.metadata.name == pod_spec.metadata.name + if randomize_name: + assert pod.metadata.name.startswith(pod_spec_name_base) + assert pod.metadata.name != pod_spec_name_base + else: + assert pod.metadata.name == pod_spec_name_base assert pod.metadata.namespace == pod_spec.metadata.namespace assert pod.spec.containers[0].image == pod_spec.spec.containers[0].image assert pod.spec.containers[0].command == pod_spec.spec.containers[0].command @@ -281,6 +296,7 @@ def test_full_pod_spec(self): name_base = "world" k = KubernetesPodOperator( task_id="task", + random_name_suffix=randomize_name, in_cluster=False, do_xcom_push=False, cluster_context="default", @@ -291,9 +307,12 @@ def test_full_pod_spec(self): ) pod = self.run_pod(k) - # make sure the kwargs takes precedence (and that name is randomized) - assert pod.metadata.name.startswith(name_base) - assert pod.metadata.name != name_base + # make sure the kwargs takes precedence (and that name is randomized when expected) + if randomize_name: + assert pod.metadata.name.startswith(name_base) + assert pod.metadata.name != name_base + else: + assert pod.metadata.name == name_base assert pod.spec.containers[0].image == image # Check labels are added from pod_template_file, the operator itself and # the pod identifying labels including Airflow version @@ -308,7 +327,8 @@ def test_full_pod_spec(self): "execution_date": mock.ANY, } - def test_pod_template_file(self): + @parameterized.expand([[True], [False]]) + def test_pod_template_file(self, randomize_name): pod_template_yaml = b""" apiVersion: v1 kind: Pod @@ -352,11 +372,16 @@ def test_pod_template_file(self): k = KubernetesPodOperator( task_id="task", + random_name_suffix=randomize_name, pod_template_file=tpl_file.name, ) pod = self.run_pod(k) - assert pod.metadata.name == "hello" + if randomize_name: + assert pod.metadata.name.startswith("hello") + assert pod.metadata.name != "hello" + else: + pod.metadata.name == "hello" # Check labels are added from pod_template_file and # the pod identifying labels including Airflow version assert pod.metadata.labels == { @@ -410,14 +435,18 @@ def test_pod_template_file(self): task_id="task", pod_template_file=tpl_file.name, name=name_base, + random_name_suffix=randomize_name, image=image, labels={"hello": "world"}, ) pod = self.run_pod(k) - # make sure the kwargs takes precedence (and that name is randomized) - assert pod.metadata.name.startswith(name_base) - assert pod.metadata.name != name_base + # make sure the kwargs takes precedence (and that name is randomized when expected) + if randomize_name: + assert pod.metadata.name.startswith(name_base) + assert pod.metadata.name != name_base + else: + assert pod.metadata.name == name_base assert pod.spec.containers[0].image == image # Check labels are added from pod_template_file, the operator itself and # the pod identifying labels including Airflow version