diff --git a/airflow/providers/docker/operators/docker_swarm.py b/airflow/providers/docker/operators/docker_swarm.py index 452d8e531a137..6b2d99fb034c8 100644 --- a/airflow/providers/docker/operators/docker_swarm.py +++ b/airflow/providers/docker/operators/docker_swarm.py @@ -22,7 +22,7 @@ import shlex from datetime import datetime from time import sleep -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal from docker import types from docker.errors import APIError @@ -107,6 +107,16 @@ class DockerSwarmOperator(DockerOperator): The resources are Resources as per the docker api [https://docker-py.readthedocs.io/en/stable/api.html#docker.types.Resources]_ This parameter has precedence on the mem_limit parameter. + :param logging_driver: The logging driver to use for container logs. Docker by default uses 'json-file'. + For more information on Docker logging drivers: https://docs.docker.com/engine/logging/configure/ + NOTE: Only drivers 'json-file' and 'gelf' are currently supported. If left empty, 'json-file' will be used. + :param logging_driver_opts: Dictionary of logging options to use with the associated logging driver chosen. + Depending on the logging driver, some options are required. + Failure to include them, will result in the operator failing. + All option values must be strings and wrapped in double quotes. + For information on 'json-file' options: https://docs.docker.com/engine/logging/drivers/json-file/ + For information on 'gelf' options: https://docs.docker.com/engine/logging/drivers/gelf/ + NOTE: 'gelf' driver requires the 'gelf-address' option to be set. """ def __init__( @@ -121,6 +131,8 @@ def __init__( networks: list[str | types.NetworkAttachmentConfig] | None = None, placement: types.Placement | list[types.Placement] | None = None, container_resources: types.Resources | None = None, + logging_driver: Literal["json-path", "gelf"] | None = None, + logging_driver_opts: dict | None = None, **kwargs, ) -> None: super().__init__(image=image, **kwargs) @@ -135,6 +147,18 @@ def __init__( self.networks = networks self.placement = placement self.container_resources = container_resources or types.Resources(mem_limit=self.mem_limit) + self.logging_driver = logging_driver + self.logging_driver_opts = logging_driver_opts + + if self.logging_driver: + supported_logging_drivers = ("json-file", "gelf") + if self.logging_driver not in supported_logging_drivers: + raise AirflowException( + f"Invalid logging driver provided: {self.logging_driver}. Must be one of: [{', '.join(supported_logging_drivers)}]" + ) + self.log_driver_config = types.DriverConfig(self.logging_driver, self.logging_driver_opts) + else: + self.log_driver_config = None def execute(self, context: Context) -> None: self.environment["AIRFLOW_TMP_DIR"] = self.tmp_dir @@ -159,6 +183,7 @@ def _run_service(self) -> None: resources=self.container_resources, networks=self.networks, placement=self.placement, + log_driver=self.log_driver_config, ), name=f"airflow-{get_random_string()}", labels={"name": f"airflow__{self.dag_id}__{self.task_id}"}, diff --git a/docs/spelling_wordlist.txt b/docs/spelling_wordlist.txt index 29d5f6d58ae71..46b441ef3e3f6 100644 --- a/docs/spelling_wordlist.txt +++ b/docs/spelling_wordlist.txt @@ -668,6 +668,8 @@ gcpcloudsql GCS gcs gdbm +Gelf +gelf generateUploadUrl Gentner geq diff --git a/tests/providers/docker/operators/test_docker_swarm.py b/tests/providers/docker/operators/test_docker_swarm.py index 4a0c6a6be7313..67976ccf4ac70 100644 --- a/tests/providers/docker/operators/test_docker_swarm.py +++ b/tests/providers/docker/operators/test_docker_swarm.py @@ -76,6 +76,8 @@ def _client_service_logs_effect(): mode=types.ServiceMode(mode="replicated", replicas=3), networks=["dummy_network"], placement=types.Placement(constraints=["node.labels.region==east"]), + logging_driver=None, + logging_driver_opts=None, ) caplog.clear() operator.execute(None) @@ -86,6 +88,7 @@ def _client_service_logs_effect(): resources=mock_obj, networks=["dummy_network"], placement=types.Placement(constraints=["node.labels.region==east"]), + log_driver=None, ) types_mock.ContainerSpec.assert_called_once_with( image="ubuntu:latest", @@ -259,6 +262,8 @@ def test_container_resources(self, types_mock, docker_api_client_patcher): cpu_reservation=100000000, mem_reservation=67108864, ), + logging_driver=None, + logging_driver_opts=None, ) operator.execute(None) @@ -273,6 +278,7 @@ def test_container_resources(self, types_mock, docker_api_client_patcher): ), networks=None, placement=None, + log_driver=None, ) types_mock.Resources.assert_not_called() @@ -355,3 +361,47 @@ def test_service_args_list(self, types_mock, docker_api_client_patcher): configs=None, secrets=None, ) + + @mock.patch("airflow.providers.docker.operators.docker_swarm.types") + def test_logging_driver(self, types_mock, docker_api_client_patcher): + mock_obj = mock.Mock() + + client_mock = mock.Mock(spec=APIClient) + client_mock.create_service.return_value = {"ID": "some_id"} + client_mock.images.return_value = [] + client_mock.pull.return_value = [b'{"status":"pull log"}'] + client_mock.tasks.return_value = [{"Status": {"State": "complete"}}] + 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 + + docker_api_client_patcher.return_value = client_mock + + operator = DockerSwarmOperator( + image="", logging_driver="json-file", task_id="unittest", enable_logging=False + ) + + assert operator.logging_driver == "json-file" + + @mock.patch("airflow.providers.docker.operators.docker_swarm.types") + def test_invalid_logging_driver(self, types_mock, docker_api_client_patcher): + mock_obj = mock.Mock() + + client_mock = mock.Mock(spec=APIClient) + client_mock.create_service.return_value = {"ID": "some_id"} + client_mock.images.return_value = [] + client_mock.pull.return_value = [b'{"status":"pull log"}'] + client_mock.tasks.return_value = [{"Status": {"State": "complete"}}] + 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 + + docker_api_client_patcher.return_value = client_mock + + msg = "Invalid logging driver provided: json. Must be one of: [json-file, gelf]" + with pytest.raises(AirflowException) as e: + # Exception is raised in __init__() + DockerSwarmOperator(image="", logging_driver="json", task_id="unittest", enable_logging=False) + assert str(e.value) == msg