Enable specifying dns and dns_search options for DockerOperator (apac…

…he#3860) Enable specifying dns and dns_search options for DockerOperator
cfei18 · Jan 23, 2019 · 39cc181 · 39cc181
1 parent 3fecb28
commit 39cc181
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py
@@ -53,6 +53,10 @@ class DockerOperator(BaseOperator):
         This value gets multiplied with 1024. See
         https://docs.docker.com/engine/reference/run/#cpu-share-constraint
     :type cpus: float
+    :param dns: Docker custom DNS servers
+    :type dns: list of strings
+    :param dns_search: Docker custom DNS search domain
+    :type dns_search: list of strings
     :param docker_url: URL of the host running the docker daemon.
         Default is unix://var/run/docker.sock
     :type docker_url: str
@@ -127,13 +131,17 @@ def __init__(
             xcom_push=False,
             xcom_all=False,
             docker_conn_id=None,
+            dns=None,
+            dns_search=None,
             *args,
             **kwargs):
 
         super(DockerOperator, self).__init__(*args, **kwargs)
         self.api_version = api_version
         self.command = command
         self.cpus = cpus
+        self.dns = dns
+        self.dns_search = dns_search
         self.docker_url = docker_url
         self.environment = environment or {}
         self.force_pull = force_pull
@@ -203,7 +211,9 @@ def execute(self, context):
                 host_config=self.cli.create_host_config(
                     binds=self.volumes,
                     network_mode=self.network_mode,
-                    shm_size=self.shm_size),
+                    shm_size=self.shm_size,
+                    dns=self.dns,
+                    dns_search=self.dns_search),
                 image=image,
                 mem_limit=self.mem_limit,
                 user=self.user,

diff --git a/tests/operators/docker_operator.py b/tests/operators/docker_operator.py
@@ -77,7 +77,9 @@ def test_execute(self, client_class_mock, mkdtemp_mock):
         client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path',
                                                                  '/mkdtemp:/tmp/airflow'],
                                                           network_mode='bridge',
-                                                          shm_size=1000)
+                                                          shm_size=1000,
+                                                          dns=None,
+                                                          dns_search=None)
         client_mock.images.assert_called_with(name='ubuntu:latest')
         client_mock.logs.assert_called_with(container='some_id', stream=True)
         client_mock.pull.assert_called_with('ubuntu:latest', stream=True)