diff --git a/providers/docker/src/airflow/providers/docker/operators/docker.py b/providers/docker/src/airflow/providers/docker/operators/docker.py index e81579a4f1c23..e87f8d08e1664 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker.py @@ -534,6 +534,15 @@ def on_kill(self) -> None: self.log.info("Not attempting to kill container as it was not created") return self.cli.stop(self.container["Id"]) + if self.auto_remove == "force": + try: + self.cli.remove_container(self.container["Id"], force=True) + except APIError as e: + self.log.info( + "Failed to remove docker container %s during on_kill; it may already be gone: %s", + self.container["Id"], + e, + ) @staticmethod def unpack_environment_variables(env_str: str) -> dict: diff --git a/providers/docker/tests/unit/docker/operators/test_docker.py b/providers/docker/tests/unit/docker/operators/test_docker.py index a753894d48f8c..81fd757ddf765 100644 --- a/providers/docker/tests/unit/docker/operators/test_docker.py +++ b/providers/docker/tests/unit/docker/operators/test_docker.py @@ -121,6 +121,52 @@ def test_on_kill_client_created(docker_api_client_patcher, container_exists): docker_api_client_patcher.return_value.stop.assert_not_called() +def test_on_kill_auto_remove_force(docker_api_client_patcher): + """Test that on_kill removes the container when auto_remove='force'.""" + op = DockerOperator( + image=TEST_IMAGE, hostname=TEST_DOCKER_URL, task_id="test_on_kill", auto_remove="force" + ) + op.container = {"Id": "some_id"} + + op.hook.get_conn() + op.on_kill() + + docker_api_client_patcher.return_value.stop.assert_called_once_with("some_id") + docker_api_client_patcher.return_value.remove_container.assert_called_once_with("some_id", force=True) + + +@pytest.mark.parametrize("auto_remove", ["success", "never"]) +def test_on_kill_auto_remove_not_force(docker_api_client_patcher, auto_remove): + """Test that on_kill does NOT remove the container when auto_remove is not 'force'.""" + op = DockerOperator( + image=TEST_IMAGE, hostname=TEST_DOCKER_URL, task_id="test_on_kill", auto_remove=auto_remove + ) + op.container = {"Id": "some_id"} + + op.hook.get_conn() + op.on_kill() + + docker_api_client_patcher.return_value.stop.assert_called_once_with("some_id") + docker_api_client_patcher.return_value.remove_container.assert_not_called() + + +def test_on_kill_auto_remove_force_api_error(docker_api_client_patcher): + """Test that on_kill logs and swallows APIError when remove_container fails.""" + docker_api_client_patcher.return_value.remove_container.side_effect = APIError("container already gone") + + op = DockerOperator( + image=TEST_IMAGE, hostname=TEST_DOCKER_URL, task_id="test_on_kill", auto_remove="force" + ) + op.container = {"Id": "some_id"} + + op.hook.get_conn() + # Must not raise — on_kill is best-effort + op.on_kill() + + docker_api_client_patcher.return_value.stop.assert_called_once_with("some_id") + docker_api_client_patcher.return_value.remove_container.assert_called_once_with("some_id", force=True) + + def test_on_kill_client_not_created(docker_api_client_patcher): """Test operator on_kill method if APIClient not created in case of error.""" docker_api_client_patcher.side_effect = APIError("Fake Client Error")