diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py index 44b72e53ea2fd..e2e85d94b23dc 100644 --- a/airflow/providers/amazon/aws/operators/ecs.py +++ b/airflow/providers/amazon/aws/operators/ecs.py @@ -86,6 +86,9 @@ class ECSOperator(BaseOperator): # pylint: disable=too-many-instance-attributes :param placement_constraints: an array of placement constraint objects to use for the task :type placement_constraints: list + :param placement_strategy: an array of placement strategy objects to use for + the task + :type placement_strategy: list :param platform_version: the platform version on which your task is running :type platform_version: str :param network_configuration: the network configuration for the task @@ -124,6 +127,7 @@ def __init__( launch_type='EC2', group=None, placement_constraints=None, + placement_strategy=None, platform_version='LATEST', network_configuration=None, tags=None, @@ -143,6 +147,7 @@ def __init__( self.launch_type = launch_type self.group = group self.placement_constraints = placement_constraints + self.placement_strategy = placement_strategy self.platform_version = platform_version self.network_configuration = network_configuration @@ -180,6 +185,8 @@ def execute(self, context): run_opts['group'] = self.group if self.placement_constraints is not None: run_opts['placementConstraints'] = self.placement_constraints + if self.placement_strategy is not None: + run_opts['placementStrategy'] = self.placement_strategy if self.network_configuration is not None: run_opts['networkConfiguration'] = self.network_configuration if self.tags is not None: diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py index 16f32b6e5c2c3..73097f185e795 100644 --- a/tests/providers/amazon/aws/operators/test_ecs.py +++ b/tests/providers/amazon/aws/operators/test_ecs.py @@ -67,6 +67,7 @@ def set_up_operator(self, aws_hook_mock, **kwargs): 'placement_constraints': [ {'expression': 'attribute:ecs.instance-type =~ t2.*', 'type': 'memberOf'} ], + 'placement_strategy': [{'field': 'memory', 'type': 'binpack'}], 'network_configuration': { 'awsvpcConfiguration': {'securityGroups': ['sg-123abc'], 'subnets': ['subnet-123456ab']} }, @@ -125,6 +126,7 @@ def test_execute_without_failures(self, launch_type, tags, check_mock, wait_mock taskDefinition='t', group='group', placementConstraints=[{'expression': 'attribute:ecs.instance-type =~ t2.*', 'type': 'memberOf'}], + placementStrategy=[{'field': 'memory', 'type': 'binpack'}], networkConfiguration={ 'awsvpcConfiguration': {'securityGroups': ['sg-123abc'], 'subnets': ['subnet-123456ab']} }, @@ -156,6 +158,7 @@ def test_execute_with_failures(self): taskDefinition='t', group='group', placementConstraints=[{'expression': 'attribute:ecs.instance-type =~ t2.*', 'type': 'memberOf'}], + placementStrategy=[{'field': 'memory', 'type': 'binpack'}], networkConfiguration={ 'awsvpcConfiguration': {'securityGroups': ['sg-123abc'], 'subnets': ['subnet-123456ab'],} },