[AIRFLOW-6574] Adding private_environment to docker operator. (#7671)
The docker operator currently does not have a means to pass in an
environment dict that is not exposed to the frontend.
- Updating docs and ensuring code is flake8.
- Adding a test and updating documentation.
diff --git a/airflow/providers/docker/operators/docker.py b/airflow/providers/docker/operators/docker.py
index b02801a..f8a8923 100644
--- a/airflow/providers/docker/operators/docker.py
+++ b/airflow/providers/docker/operators/docker.py
@@ -65,6 +65,9 @@
:type docker_url: str
:param environment: Environment variables to set in the container. (templated)
:type environment: dict
+ :param private_environment: Private environment variables to set in the container.
+ These are not templated, and hidden from the website.
+ :type private_environment: dict
: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.
@@ -136,6 +139,7 @@
cpus: float = 1.0,
docker_url: str = 'unix://var/run/docker.sock',
environment: Optional[Dict] = None,
+ private_environment: Optional[Dict] = None,
force_pull: bool = False,
mem_limit: Optional[Union[float, str]] = None,
host_tmp_dir: Optional[str] = None,
@@ -169,6 +173,7 @@
self.dns_search = dns_search
self.docker_url = docker_url
self.environment = environment or {}
+ self._private_environment = private_environment or {}
self.force_pull = force_pull
self.image = image
self.mem_limit = mem_limit
@@ -218,7 +223,7 @@
self.container = self.cli.create_container(
command=self.get_command(),
name=self.container_name,
- environment=self.environment,
+ environment={**self.environment, **self._private_environment},
host_config=self.cli.create_host_config(
auto_remove=self.auto_remove,
binds=self.volumes,
diff --git a/tests/providers/docker/operators/test_docker.py b/tests/providers/docker/operators/test_docker.py
index 068eaef..44aad34 100644
--- a/tests/providers/docker/operators/test_docker.py
+++ b/tests/providers/docker/operators/test_docker.py
@@ -50,8 +50,9 @@
client_class_mock.return_value = client_mock
operator = DockerOperator(api_version='1.19', command='env', environment={'UNIT': 'TEST'},
- image='ubuntu:latest', network_mode='bridge', owner='unittest',
- task_id='unittest', volumes=['/host/path:/container/path'],
+ private_environment={'PRIVATE': 'MESSAGE'}, 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', container_name='test_container',
tty=True)
@@ -64,7 +65,8 @@
name='test_container',
environment={
'AIRFLOW_TMP_DIR': '/tmp/airflow',
- 'UNIT': 'TEST'
+ 'UNIT': 'TEST',
+ 'PRIVATE': 'MESSAGE'
},
host_config=host_config,
image='ubuntu:latest',
@@ -89,6 +91,15 @@
decode=True)
client_mock.wait.assert_called_once_with('some_id')
+ def test_private_environment_is_private(self):
+ operator = DockerOperator(private_environment={'PRIVATE': 'MESSAGE'},
+ image='ubuntu:latest',
+ task_id='unittest')
+ self.assertEqual(
+ operator._private_environment, {'PRIVATE': 'MESSAGE'},
+ "To keep this private, it must be an underscored attribute."
+ )
+
@mock.patch('airflow.providers.docker.operators.docker.tls.TLSConfig')
@mock.patch('airflow.providers.docker.operators.docker.APIClient')
def test_execute_tls(self, client_class_mock, tls_class_mock):
@@ -259,6 +270,7 @@
'api_version': '1.19',
'command': 'env',
'environment': {'UNIT': 'TEST'},
+ 'private_environment': {'PRIVATE': 'MESSAGE'},
'image': 'ubuntu:latest',
'network_mode': 'bridge',
'owner': 'unittest',
