AIRFLOW-5854: Add support for tty parameter in Docker related opera…

…tors (#6542) (cherry-picked from commit 7ab9e95)
apache · Nov 11, 2019 · 10534a1 · 10534a1
1 parent d3899dc
commit 10534a1
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py
@@ -120,6 +120,9 @@ class DockerOperator(BaseOperator):
     :param shm_size: Size of ``/dev/shm`` in bytes. The size must be
         greater than 0. If omitted uses system default.
     :type shm_size: int
+    :param tty: Allocate pseudo-TTY to the container
+        This needs to be set see logs of the Docker container.
+    :type tty: bool
     """
     template_fields = ('command', 'environment', 'container_name')
     template_ext = ('.sh', '.bash',)
@@ -154,6 +157,7 @@ def __init__(
             dns_search=None,
             auto_remove=False,
             shm_size=None,
+            tty=False,
             *args,
             **kwargs):
 
@@ -185,6 +189,7 @@ def __init__(
         self.xcom_all = xcom_all
         self.docker_conn_id = docker_conn_id
         self.shm_size = shm_size
+        self.tty = tty
 
         self.cli = None
         self.container = None
@@ -237,7 +242,8 @@ def execute(self, context):
                     mem_limit=self.mem_limit),
                 image=self.image,
                 user=self.user,
-                working_dir=self.working_dir
+                working_dir=self.working_dir,
+                tty=self.tty,
             )
             self.cli.start(self.container['Id'])
 

diff --git a/tests/operators/test_docker_operator.py b/tests/operators/test_docker_operator.py
@@ -53,7 +53,8 @@ def test_execute(self, client_class_mock, mkdtemp_mock):
                                   image='ubuntu:latest', network_mode='bridge', owner='unittest',
                                   task_id='unittest', volumes=['/host/path:/container/path'],
                                   working_dir='/container/path', shm_size=1000,
-                                  host_tmp_dir='/host/airflow', container_name='test_container')
+                                  host_tmp_dir='/host/airflow', container_name='test_container',
+                                  tty=True)
         operator.execute(None)
 
         client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None,
@@ -68,7 +69,8 @@ def test_execute(self, client_class_mock, mkdtemp_mock):
                                                         host_config=host_config,
                                                         image='ubuntu:latest',
                                                         user=None,
-                                                        working_dir='/container/path'
+                                                        working_dir='/container/path',
+                                                        tty=True
                                                         )
         client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path',
                                                                  '/mkdtemp:/tmp/airflow'],