Added logging device and logging device options

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}"},

docs/spelling_wordlist.txt

@@ -668,6 +668,8 @@ gcpcloudsql
 GCS
 gcs
 gdbm
+Gelf
+gelf
 generateUploadUrl
 Gentner
 geq

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