From 1d18d28b36ced9742193539ccc5e318120f7e26d Mon Sep 17 00:00:00 2001 From: RachitSharma2001 Date: Mon, 1 May 2023 08:41:49 -0700 Subject: [PATCH 1/3] Add options to KubernetesPodOperator --- .../cncf/kubernetes/operators/pod.py | 9 +++++ .../cncf/kubernetes/operators/test_pod.py | 35 +++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 6f5a3db603828..29f8c985a195e 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -282,6 +282,9 @@ def __init__( security_context: dict | None = None, container_security_context: dict | None = None, dnspolicy: str | None = None, + dns_config: dict | None = None, + hostname: str | None = None, + subdomain: str | None = None, schedulername: str | None = None, full_pod_spec: k8s.V1Pod | None = None, init_containers: list[k8s.V1Container] | None = None, @@ -351,6 +354,9 @@ def __init__( self.security_context = security_context or {} self.container_security_context = container_security_context self.dnspolicy = dnspolicy + self.dns_config = dns_config + self.hostname = hostname + self.subdomain = subdomain self.schedulername = schedulername self.full_pod_spec = full_pod_spec self.init_containers = init_containers or [] @@ -812,8 +818,11 @@ def build_pod_request_obj(self, context: Context | None = None) -> k8s.V1Pod: image_pull_secrets=self.image_pull_secrets, service_account_name=self.service_account_name, host_network=self.hostnetwork, + hostname=self.hostname, + subdomain=self.subdomain, security_context=self.security_context, dns_policy=self.dnspolicy, + dns_config=self.dns_config, scheduler_name=self.schedulername, restart_policy="Never", priority_class_name=self.priority_class_name, diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index f3fd713f9b662..a60862e89ede0 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -324,6 +324,41 @@ def test_find_pod_labels(self): "already_checked!=True,!airflow-worker" ) + @patch(HOOK_CLASS, new=MagicMock) + def test_pod_additional_options(self): + dns_config = k8s.V1PodDNSConfig( + nameservers=["192.0.2.1", "192.0.2.3"], + searches=["ns1.svc.cluster-domain.example", "my.dns.search.suffix"], + options=[ + k8s.V1PodDNSConfigOption( + name="ndots", + value="2", + ) + ], + ) + hostname = "busybox-2" + subdomain = "busybox-subdomain" + + k = KubernetesPodOperator( + namespace="default", + image="ubuntu:16.04", + cmds=["bash", "-cx"], + labels={"foo": "bar"}, + name="test", + task_id="task", + in_cluster=False, + do_xcom_push=False, + dns_config=dns_config, + hostname=hostname, + subdomain=subdomain, + ) + + self.run_pod(k) + pod_spec = k.pod.spec + assert pod_spec.dns_config == dns_config + assert pod_spec.subdomain == subdomain + assert pod_spec.hostname == hostname + @pytest.mark.parametrize( "val", [ From f1f79d3891aeacacb43e765461bf05e20203b725 Mon Sep 17 00:00:00 2001 From: RachitSharma2001 Date: Sat, 6 May 2023 10:28:10 -0700 Subject: [PATCH 2/3] Add name and type changes --- airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- tests/providers/cncf/kubernetes/operators/test_pod.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 29f8c985a195e..e30fed815b19b 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -282,7 +282,7 @@ def __init__( security_context: dict | None = None, container_security_context: dict | None = None, dnspolicy: str | None = None, - dns_config: dict | None = None, + dns_config: k8s.V1PodDNSConfig | None = None, hostname: str | None = None, subdomain: str | None = None, schedulername: str | None = None, diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index a60862e89ede0..c220b344291a1 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -325,7 +325,7 @@ def test_find_pod_labels(self): ) @patch(HOOK_CLASS, new=MagicMock) - def test_pod_additional_options(self): + def test_pod_dns_options(self): dns_config = k8s.V1PodDNSConfig( nameservers=["192.0.2.1", "192.0.2.3"], searches=["ns1.svc.cluster-domain.example", "my.dns.search.suffix"], From 1f3d3f87662709bb88a76be10e81bf7fe10a0d23 Mon Sep 17 00:00:00 2001 From: RachitSharma2001 Date: Sun, 7 May 2023 12:08:17 -0700 Subject: [PATCH 3/3] Add parameters to docstring --- airflow/providers/cncf/kubernetes/operators/pod.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index e30fed815b19b..08791b4ee15b0 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -197,6 +197,9 @@ class KubernetesPodOperator(BaseOperator): :param security_context: security options the pod should run with (PodSecurityContext). :param container_security_context: security options the container should run with. :param dnspolicy: dnspolicy for the pod. + :param dns_config: dns configuration (ip addresses, searches, options) for the pod. + :param hostname: hostname for the pod. + :param subdomain: subdomain for the pod. :param schedulername: Specify a schedulername for the pod :param full_pod_spec: The complete podSpec :param init_containers: init container for the launched Pod