diff --git a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py index 6a77f73b81055..2c403802f2f40 100644 --- a/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py +++ b/providers/docker/src/airflow/providers/docker/operators/docker_swarm.py @@ -113,6 +113,8 @@ 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 service_prefix: Prefix for the service name. The service name will be generated as + ``{service_prefix}-{random_string}``. Default is 'airflow'. :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. @@ -137,6 +139,7 @@ def __init__( networks: list[str | types.NetworkAttachmentConfig] | None = None, placement: types.Placement | list[types.Placement] | None = None, container_resources: types.Resources | None = None, + service_prefix: str = "airflow", logging_driver: Literal["json-path", "gelf"] | None = None, logging_driver_opts: dict | None = None, **kwargs, @@ -153,6 +156,7 @@ def __init__( self.networks = networks self.placement = placement self.container_resources = container_resources or types.Resources(mem_limit=self.mem_limit) + self.service_prefix = service_prefix self.logging_driver = logging_driver self.logging_driver_opts = logging_driver_opts @@ -191,7 +195,7 @@ def _run_service(self) -> None: placement=self.placement, log_driver=self.log_driver_config, ), - name=f"airflow-{get_random_string()}", + name=f"{self.service_prefix}-{get_random_string()}", labels={"name": f"airflow__{self.dag_id}__{self.task_id}"}, mode=self.mode, ) diff --git a/providers/docker/tests/unit/docker/operators/test_docker_swarm.py b/providers/docker/tests/unit/docker/operators/test_docker_swarm.py index 0887785739b1e..5b4190ec5d4bf 100644 --- a/providers/docker/tests/unit/docker/operators/test_docker_swarm.py +++ b/providers/docker/tests/unit/docker/operators/test_docker_swarm.py @@ -124,7 +124,6 @@ def _client_service_logs_effect(): assert len(csargs) == 1, "create_service called with different number of arguments than expected" assert csargs == (mock_obj,) assert cskwargs["labels"] == {"name": "airflow__adhoc_airflow__unittest"} - assert cskwargs["name"].startswith("airflow-") assert cskwargs["mode"] == types.ServiceMode(mode="replicated", replicas=3) assert client_mock.tasks.call_count == 8 client_mock.remove_service.assert_called_once_with("some_id") @@ -439,3 +438,47 @@ def test_invalid_logging_driver(self, types_mock, docker_api_client_patcher): # Exception is raised in __init__() DockerSwarmOperator(image="", logging_driver="json", task_id="unittest", enable_logging=False) assert str(e.value) == msg + + @mock.patch("airflow.providers.docker.operators.docker_swarm.types") + @pytest.mark.parametrize( + "service_prefix,expected_prefix", + [ + (None, "airflow-"), # Default case + ("airflow", "airflow-"), + ("custom-prefix", "custom-prefix-"), + ("prefix_with_underscores", "prefix_with_underscores-"), + ], + ) + def test_service_prefix(self, types_mock, docker_api_client_patcher, service_prefix, expected_prefix): + """Test that service prefix is correctly applied in service creation.""" + 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", "ContainerStatus": {"ContainerID": "some_id"}}} + ] + 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_kwargs = {} + if service_prefix: + operator_kwargs["service_prefix"] = service_prefix + + operator = DockerSwarmOperator( + image="ubuntu:latest", + task_id="unittest", + auto_remove="success", + enable_logging=False, + **operator_kwargs, + ) + operator.execute(None) + + _, cskwargs = client_mock.create_service.call_args_list[0] + assert cskwargs["name"].startswith(expected_prefix)