diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py index e1acb5799405bb..a8d298fe391195 100644 --- a/airflow/operators/docker_operator.py +++ b/airflow/operators/docker_operator.py @@ -57,6 +57,8 @@ class DockerOperator(BaseOperator): :type auto_remove: bool :param command: Command to be run in the container. (templated) :type command: str or list + :param container_name: Name of the container. + :type container_name: str :param cpus: Number of CPUs to assign to the container. This value gets multiplied with 1024. See https://docs.docker.com/engine/reference/run/#cpu-share-constraint @@ -128,6 +130,7 @@ def __init__( image, api_version=None, command=None, + container_name=None, cpus=1.0, docker_url='unix://var/run/docker.sock', environment=None, @@ -158,6 +161,7 @@ def __init__( self.api_version = api_version self.auto_remove = auto_remove self.command = command + self.container_name = container_name self.cpus = cpus self.dns = dns self.dns_search = dns_search @@ -220,6 +224,7 @@ def execute(self, context): self.container = self.cli.create_container( command=self.get_command(), + name=self.container_name, environment=self.environment, host_config=self.cli.create_host_config( auto_remove=self.auto_remove, diff --git a/tests/operators/test_docker_operator.py b/tests/operators/test_docker_operator.py index b60f069d0d87e0..17db0e4929a5ce 100644 --- a/tests/operators/test_docker_operator.py +++ b/tests/operators/test_docker_operator.py @@ -53,13 +53,14 @@ def test_execute(self, client_class_mock, mkdtemp_mock): image='ubuntu:latest', network_mode='bridge', owner='unittest', task_id='unittest', volumes=['/host/path:/container/path'], working_dir='/container/path', shm_size=1000, - host_tmp_dir='/host/airflow') + host_tmp_dir='/host/airflow', container_name='test_container') operator.execute(None) client_class_mock.assert_called_with(base_url='unix://var/run/docker.sock', tls=None, version='1.19') client_mock.create_container.assert_called_with(command='env', + name='test_container', environment={ 'AIRFLOW_TMP_DIR': '/tmp/airflow', 'UNIT': 'TEST'