diff --git a/airflow/providers/docker/operators/docker_swarm.py b/airflow/providers/docker/operators/docker_swarm.py index 2d5373c840f17..a1f3f0b2bcf2d 100644 --- a/airflow/providers/docker/operators/docker_swarm.py +++ b/airflow/providers/docker/operators/docker_swarm.py @@ -17,7 +17,6 @@ """Run ephemeral Docker Swarm services""" from typing import List, Optional, Union -import requests from docker import types from airflow.exceptions import AirflowException @@ -204,12 +203,6 @@ def _stream_logs_to_output(self) -> None: while True: try: log = next(logs) - # TODO: Remove this clause once https://github.com/docker/docker-py/issues/931 is fixed - except requests.exceptions.ConnectionError: - # If the service log stream stopped sending messages, check if it the service has - # terminated. - if self._has_service_terminated(): - break except StopIteration: # If the service log stream terminated, stop fetching logs further. break diff --git a/setup.py b/setup.py index 0374e34bb6332..4f2db05c3eb40 100644 --- a/setup.py +++ b/setup.py @@ -260,7 +260,7 @@ def write_version(filename: str = os.path.join(*[my_dir, "airflow", "git_version 'sphinxcontrib-spelling==7.2.1', ] docker = [ - 'docker', + 'docker>=5.0.3', ] drill = ['sqlalchemy-drill>=1.1.0', 'sqlparse>=0.4.1'] druid = [ diff --git a/tests/providers/docker/operators/test_docker_swarm.py b/tests/providers/docker/operators/test_docker_swarm.py index 8523644888de4..09207b425dbb8 100644 --- a/tests/providers/docker/operators/test_docker_swarm.py +++ b/tests/providers/docker/operators/test_docker_swarm.py @@ -20,7 +20,6 @@ from unittest import mock import pytest -import requests from docker import APIClient, types from parameterized import parameterized @@ -184,53 +183,6 @@ def test_non_complete_service_raises_error(self, status, types_mock, client_clas operator.execute(None) assert str(ctx.value) == msg - @mock.patch('airflow.providers.docker.operators.docker.APIClient') - @mock.patch('airflow.providers.docker.operators.docker_swarm.types') - def test_logging_with_requests_timeout(self, types_mock, client_class_mock): - - mock_obj = mock.Mock() - - def _client_tasks_side_effect(): - for _ in range(2): - yield [{'Status': {'State': 'pending'}}] - while True: - yield [{'Status': {'State': 'complete'}}] - - def _client_service_logs_effect(): - yield b'Testing is awesome.' - raise requests.exceptions.ConnectionError('') - - client_mock = mock.Mock(spec=APIClient) - client_mock.create_service.return_value = {'ID': 'some_id'} - client_mock.service_logs.return_value = _client_service_logs_effect() - client_mock.images.return_value = [] - client_mock.pull.return_value = [b'{"status":"pull log"}'] - client_mock.tasks.side_effect = _client_tasks_side_effect() - types_mock.TaskTemplate.return_value = mock_obj - types_mock.ContainerSpec.return_value = mock_obj - types_mock.RestartPolicy.return_value = mock_obj - types_mock.Resources.return_value = mock_obj - - client_class_mock.return_value = client_mock - - operator = DockerSwarmOperator( - api_version='1.19', - command='env', - environment={'UNIT': 'TEST'}, - image='ubuntu:latest', - mem_limit='128m', - user='unittest', - task_id='unittest', - auto_remove=True, - tty=True, - enable_logging=True, - ) - operator.execute(None) - - client_mock.service_logs.assert_called_once_with( - 'some_id', follow=True, stdout=True, stderr=True, is_tty=True - ) - def test_on_kill(self): client_mock = mock.Mock(spec=APIClient)