From d590e5e7679322bebb1472fa8c7ec6d183e4154a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Ponte?= Date: Mon, 11 May 2020 19:48:13 +0200 Subject: [PATCH] Add option to propagate tags in ECSOperator (#8811) Co-authored-by: Joao Ponte --- airflow/providers/amazon/aws/operators/ecs.py | 5 ++++- tests/providers/amazon/aws/operators/test_ecs.py | 7 +++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py index 51f5bab4b9..ab729c52e8 100644 --- a/airflow/providers/amazon/aws/operators/ecs.py +++ b/airflow/providers/amazon/aws/operators/ecs.py @@ -113,7 +113,7 @@ class ECSOperator(BaseOperator): # pylint: disable=too-many-instance-attributes aws_conn_id=None, region_name=None, launch_type='EC2', group=None, placement_constraints=None, platform_version='LATEST', network_configuration=None, tags=None, awslogs_group=None, - awslogs_region=None, awslogs_stream_prefix=None, **kwargs): + awslogs_region=None, awslogs_stream_prefix=None, propagate_tags=None, **kwargs): super().__init__(**kwargs) self.aws_conn_id = aws_conn_id @@ -131,6 +131,7 @@ class ECSOperator(BaseOperator): # pylint: disable=too-many-instance-attributes self.awslogs_group = awslogs_group self.awslogs_stream_prefix = awslogs_stream_prefix self.awslogs_region = awslogs_region + self.propagate_tags = propagate_tags if self.awslogs_region is None: self.awslogs_region = region_name @@ -165,6 +166,8 @@ class ECSOperator(BaseOperator): # pylint: disable=too-many-instance-attributes run_opts['networkConfiguration'] = self.network_configuration if self.tags is not None: run_opts['tags'] = [{'key': k, 'value': v} for (k, v) in self.tags.items()] + if self.propagate_tags is not None: + run_opts['propagateTags'] = self.propagate_tags response = self.client.run_task(**run_opts) diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py index 875abf1c7a..49b27d2ea7 100644 --- a/tests/providers/amazon/aws/operators/test_ecs.py +++ b/tests/providers/amazon/aws/operators/test_ecs.py @@ -72,7 +72,8 @@ class TestECSOperator(unittest.TestCase): 'securityGroups': ['sg-123abc'], 'subnets': ['subnet-123456ab'] } - } + }, + 'propagate_tags': 'TASK_DEFINITION' } self.ecs = ECSOperator(**self.ecs_operator_args, **kwargs) self.ecs.get_hook() @@ -137,6 +138,7 @@ class TestECSOperator(unittest.TestCase): 'subnets': ['subnet-123456ab'] } }, + propagateTags='TASK_DEFINITION', **extend_args ) @@ -173,7 +175,8 @@ class TestECSOperator(unittest.TestCase): 'securityGroups': ['sg-123abc'], 'subnets': ['subnet-123456ab'], } - } + }, + propagateTags='TASK_DEFINITION' ) def test_wait_end_tasks(self):