Add placement_strategy option (#9444)

apache · Aug 31, 2020 · f40ac9b · f40ac9b
1 parent 901ecf7
commit f40ac9b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git 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
@@ -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'],}
             },