[airflow] branch master updated: Remove failed DockerOperator tasks with auto_remove=True (#13532) (#13993)

[po...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

potiuk pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/master by this push:
     new ba54afe  Remove failed DockerOperator tasks with auto_remove=True (#13532) (#13993)
ba54afe is described below

commit ba54afe58b7cbd3711aca23252027fbd034cca41
Author: Danilo Trombino <da...@gmail.com>
AuthorDate: Sun Jan 31 20:23:45 2021 +0100

    Remove failed DockerOperator tasks with auto_remove=True (#13532) (#13993)
    
    * Remove failed DockerOperator tasks with auto_remove=True
    
    Removes exited containers if a task based on DockerOperator fails
    with StatusCode!=0 and auto_remove=True
---
 airflow/providers/docker/operators/docker.py    | 2 ++
 tests/providers/docker/operators/test_docker.py | 9 +++++++++
 2 files changed, 11 insertions(+)

diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py
index 0440d21..8662d4d 100644
--- a/airflow/providers/docker/operators/docker.py
+++ b/airflow/providers/docker/operators/docker.py
@@ -264,6 +264,8 @@ class DockerOperator(BaseOperator):
 
             result = self.cli.wait(self.container['Id'])
             if result['StatusCode'] != 0:
+                if self.auto_remove:
+                    self.cli.remove_container(self.container['Id'])
                 raise AirflowException('docker container failed: ' + repr(result))
 
             # duplicated conditional logic because of expensive operation
diff --git a/tests/providers/docker/operators/test_docker.py b/tests/providers/docker/operators/test_docker.py
index 0a2f838..58fc0f5 100644
--- a/tests/providers/docker/operators/test_docker.py
+++ b/tests/providers/docker/operators/test_docker.py
@@ -167,6 +167,15 @@ class TestDockerOperator(unittest.TestCase):
         with pytest.raises(AirflowException):
             operator.execute(None)
 
+    def test_auto_remove_container_fails(self):
+        self.client_mock.wait.return_value = {"StatusCode": 1}
+        operator = DockerOperator(image='ubuntu', owner='unittest', task_id='unittest', auto_remove=True)
+        operator.container = {'Id': 'some_id'}
+        with pytest.raises(AirflowException):
+            operator.execute(None)
+
+        self.client_mock.remove_container.assert_called_once_with('some_id')
+
     @staticmethod
     def test_on_kill():
         client_mock = mock.Mock(spec=APIClient)