Skip to content

Commit

Permalink
[AIRFLOW-2891] allow configurable docker_operator container name (apa…
Browse files Browse the repository at this point in the history
  • Loading branch information
rocketbitz authored and BasPH committed Jul 31, 2019
1 parent f76d9da commit c3f01f8
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 1 deletion.
5 changes: 5 additions & 0 deletions airflow/operators/docker_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,8 @@ class DockerOperator(BaseOperator):
:type api_version: str
: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
Expand Down Expand Up @@ -123,6 +125,7 @@ def __init__(
image: str,
api_version: str = None,
command: Union[str, List[str]] = None,
container_name: str = None,
cpus: float = 1.0,
docker_url: str = 'unix://var/run/docker.sock',
environment: Dict = None,
Expand Down Expand Up @@ -152,6 +155,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
Expand Down Expand Up @@ -215,6 +219,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,
Expand Down
3 changes: 2 additions & 1 deletion tests/operators/test_docker_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down

0 comments on commit c3f01f8

Please sign in to comment.