diff --git a/airflow/operators/docker_operator.py b/airflow/operators/docker_operator.py index 517199be513cb..07ad5f5c74d38 100644 --- a/airflow/operators/docker_operator.py +++ b/airflow/operators/docker_operator.py @@ -43,6 +43,7 @@ class DockerOperator(BaseOperator): be provided with the parameter ``docker_conn_id``. :param image: Docker image from which to create the container. + If image tag is omitted, "latest" will be used. :type image: str :param api_version: Remote API version. Set to ``auto`` to automatically detect the server's version. @@ -62,7 +63,7 @@ class DockerOperator(BaseOperator): :type docker_url: str :param environment: Environment variables to set in the container. (templated) :type environment: dict - :param force_pull: Pull the docker image on every run. Default is false. + :param force_pull: Pull the docker image on every run. Default is False. :type force_pull: bool :param mem_limit: Maximum amount of memory the container can use. Either a float value, which represents the limit in bytes, @@ -187,35 +188,28 @@ def execute(self, context): tls=tls_config ) - if ':' not in self.image: - image = self.image + ':latest' - else: - image = self.image - - if self.force_pull or len(self.cli.images(name=image)) == 0: - self.log.info('Pulling docker image %s', image) - for l in self.cli.pull(image, stream=True): + if self.force_pull or len(self.cli.images(name=self.image)) == 0: + self.log.info('Pulling docker image %s', self.image) + for l in self.cli.pull(self.image, stream=True): output = json.loads(l.decode('utf-8')) self.log.info("%s", output['status']) - cpu_shares = int(round(self.cpus * 1024)) - with TemporaryDirectory(prefix='airflowtmp') as host_tmp_dir: self.environment['AIRFLOW_TMP_DIR'] = self.tmp_dir self.volumes.append('{0}:{1}'.format(host_tmp_dir, self.tmp_dir)) self.container = self.cli.create_container( command=self.get_command(), - cpu_shares=cpu_shares, environment=self.environment, host_config=self.cli.create_host_config( binds=self.volumes, network_mode=self.network_mode, shm_size=self.shm_size, dns=self.dns, - dns_search=self.dns_search), - image=image, - mem_limit=self.mem_limit, + dns_search=self.dns_search, + cpu_shares=int(round(self.cpus * 1024)), + mem_limit=self.mem_limit), + image=self.image, user=self.user, working_dir=self.working_dir ) diff --git a/tests/operators/docker_operator.py b/tests/operators/test_docker_operator.py similarity index 97% rename from tests/operators/docker_operator.py rename to tests/operators/test_docker_operator.py index ea90c53c28d82..a7d63e4ebcf1d 100644 --- a/tests/operators/docker_operator.py +++ b/tests/operators/test_docker_operator.py @@ -64,20 +64,22 @@ def test_execute(self, client_class_mock, mkdtemp_mock): 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', cpu_shares=1024, + client_mock.create_container.assert_called_with(command='env', environment={ 'AIRFLOW_TMP_DIR': '/tmp/airflow', 'UNIT': 'TEST' }, host_config=host_config, image='ubuntu:latest', - mem_limit=None, user=None, + user=None, working_dir='/container/path' ) client_mock.create_host_config.assert_called_with(binds=['/host/path:/container/path', '/mkdtemp:/tmp/airflow'], network_mode='bridge', shm_size=1000, + cpu_shares=1024, + mem_limit=None, dns=None, dns_search=None) client_mock.images.assert_called_with(name='ubuntu:latest') diff --git a/tests/operators/hive_operator.py b/tests/operators/test_hive_operator.py similarity index 100% rename from tests/operators/hive_operator.py rename to tests/operators/test_hive_operator.py diff --git a/tests/operators/latest_only_operator.py b/tests/operators/test_latest_only_operator.py similarity index 100% rename from tests/operators/latest_only_operator.py rename to tests/operators/test_latest_only_operator.py diff --git a/tests/operators/python_operator.py b/tests/operators/test_python_operator.py similarity index 100% rename from tests/operators/python_operator.py rename to tests/operators/test_python_operator.py diff --git a/tests/operators/s3_to_hive_operator.py b/tests/operators/test_s3_to_hive_operator.py similarity index 100% rename from tests/operators/s3_to_hive_operator.py rename to tests/operators/test_s3_to_hive_operator.py diff --git a/tests/operators/slack_operator.py b/tests/operators/test_slack_operator.py similarity index 100% rename from tests/operators/slack_operator.py rename to tests/operators/test_slack_operator.py diff --git a/tests/operators/subdag_operator.py b/tests/operators/test_subdag_operator.py similarity index 100% rename from tests/operators/subdag_operator.py rename to tests/operators/test_subdag_operator.py