DynamoDB: Allow removal of multiple listitems (#5767)

moto/dynamodb/models/__init__.py

@@ -1588,6 +1588,7 @@ class DynamoDBBackend(BaseBackend):
                 table=table,
             )
             validated_ast = validator.validate()
+            validated_ast.normalize()
             try:
                 UpdateExpressionExecutor(
                     validated_ast, item, expression_attribute_names

moto/dynamodb/parsing/ast_nodes.py

@@ -23,10 +23,10 @@ class Node(metaclass=abc.ABCMeta):
 
     def validate(self):
         if self.type == "UpdateExpression":
-            nr_of_clauses = len(self.find_clauses(UpdateExpressionAddClause))
+            nr_of_clauses = len(self.find_clauses([UpdateExpressionAddClause]))
             if nr_of_clauses > 1:
                 raise TooManyAddClauses()
-            set_actions = self.find_clauses(UpdateExpressionSetAction)
+            set_actions = self.find_clauses([UpdateExpressionSetAction])
             # set_attributes = ["attr", "map.attr", attr.list[2], ..]
             set_attributes = [s.children[0].to_str() for s in set_actions]
             # We currently only check for duplicates
@@ -34,13 +34,53 @@ class Node(metaclass=abc.ABCMeta):
             if len(set_attributes) != len(set(set_attributes)):
                 raise DuplicateUpdateExpression(set_attributes)
 
-    def find_clauses(self, clause_type):
+    def normalize(self):
+        """
+        Flatten the Add-/Delete-/Remove-/Set-Action children within this Node
+        """
+        if self.type == "UpdateExpression":
+            # We can have multiple REMOVE attr[idx] expressions, such as attr[i] and attr[i+2]
+            # If we remove attr[i] first, attr[i+2] suddenly refers to a different item
+            # So we sort them in reverse order - we can remove attr[i+2] first, attr[i] still refers to the same item
+
+            # Behaviour that is unknown, for now:
+            # What happens if we SET and REMOVE on the same list - what takes precedence?
+            # We're assuming this is executed in original order
+
+            remove_actions = []
+            sorted_actions = []
+            possible_clauses = [
+                UpdateExpressionAddAction,
+                UpdateExpressionDeleteAction,
+                UpdateExpressionRemoveAction,
+                UpdateExpressionSetAction,
+            ]
+            for action in self.find_clauses(possible_clauses):
+                if isinstance(action, UpdateExpressionRemoveAction):
+                    # Keep these separate for now
+                    remove_actions.append(action)
+                else:
+                    if len(remove_actions) > 0:
+                        # Remove-actions were found earlier
+                        # Now that we have other action-types, that means we've found all possible Remove-actions
+                        # Sort them appropriately
+                        sorted_actions.extend(sorted(remove_actions, reverse=True))
+                        remove_actions.clear()
+                    # Add other actions by insertion order
+                    sorted_actions.append(action)
+            # Remove actions were found last
+            if len(remove_actions) > 0:
+                sorted_actions.extend(sorted(remove_actions, reverse=True))
+
+            self.children = sorted_actions
+
+    def find_clauses(self, clause_types):
         clauses = []
         for child in self.children or []:
-            if isinstance(child, clause_type):
+            if type(child) in clause_types:
                 clauses.append(child)
             elif isinstance(child, Expression):
-                clauses.extend(child.find_clauses(clause_type))
+                clauses.extend(child.find_clauses(clause_types))
         return clauses
 
 
@@ -115,6 +155,16 @@ class UpdateExpressionRemoveAction(UpdateExpressionClause):
     RemoveAction => Path
     """
 
+    def _get_value(self):
+        expression_path = self.children[0]
+        expression_selector = expression_path.children[-1]
+        return expression_selector.children[0]
+
+    def __lt__(self, other):
+        self_value = self._get_value()
+
+        return self_value < other._get_value()
+
 
 class UpdateExpressionAddActions(UpdateExpressionClause):
     """

moto/dynamodb/parsing/executors.py

@@ -273,6 +273,8 @@ class UpdateExpressionExecutor(object):
         and process the nodes 1-by-1. If no specific execution for the node type is defined we can execute the children
         in order since it will be a container node that is expandable and left child will be first in the statement.
 
+        Note that, if `normalize()` is called before, the list of children will be flattened and sorted (if appropriate).
+
         Args:
             node(Node):
 

tests/test_dynamodb/test_dynamodb.py

@@ -2721,6 +2721,29 @@ def test_remove_list_index__remove_existing_index():
     result["itemlist"].should.equal({"L": [{"S": "bar1"}, {"S": "bar3"}]})
 
 
+@mock_dynamodb
+def test_remove_list_index__remove_multiple_indexes():
+    table_name = "remove-test"
+    create_table_with_list(table_name)
+    dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
+
+    table = dynamodb.Table(table_name)
+    table.put_item(
+        Item={
+            "id": "woop",
+            "bla": ["1", "2", "3", "4", "5"],
+        },
+    )
+
+    table.update_item(
+        Key={"id": "woop"}, UpdateExpression="REMOVE bla[0], bla[1], bla[2]"
+    )
+
+    result = table.get_item(Key={"id": "woop"})
+    item = result["Item"]
+    assert item["bla"] == ["4", "5"]
+
+
 @mock_dynamodb
 def test_remove_list_index__remove_existing_nested_index():
     table_name = "test_list_index_access"

tests/test_dynamodb/test_dynamodb_executor.py

@@ -2,6 +2,13 @@ import pytest
 
 from moto.dynamodb.exceptions import IncorrectOperandType, IncorrectDataType
 from moto.dynamodb.models import Item, DynamoType
+from moto.dynamodb.parsing.ast_nodes import (
+    UpdateExpression,
+    UpdateExpressionAddClause,
+    UpdateExpressionAddAction,
+    UpdateExpressionRemoveAction,
+    UpdateExpressionSetAction,
+)
 from moto.dynamodb.parsing.executors import UpdateExpressionExecutor
 from moto.dynamodb.parsing.expressions import UpdateExpressionParser
 from moto.dynamodb.parsing.validators import UpdateExpressionValidator
@@ -430,3 +437,66 @@ def test_execution_of_delete_element_from_a_string_attribute(table):
         assert False, "Must raise exception"
     except IncorrectDataType:
         assert True
+
+
+def test_normalize_with_one_action(table):
+    update_expression = "ADD s :value"
+    update_expression_ast = UpdateExpressionParser.make(update_expression)
+    item = Item(
+        hash_key=DynamoType({"S": "id"}),
+        range_key=None,
+        attrs={"id": {"S": "foo2"}, "s": {"SS": ["value1", "value2", "value3"]}},
+    )
+    validated_ast = UpdateExpressionValidator(
+        update_expression_ast,
+        expression_attribute_names=None,
+        expression_attribute_values={":value": {"SS": ["value2", "value5"]}},
+        item=item,
+        table=table,
+    ).validate()
+    validated_ast.children.should.have.length_of(1)
+    validated_ast.children[0].should.be.a(UpdateExpressionAddClause)
+
+    validated_ast.normalize()
+    validated_ast.children.should.have.length_of(1)
+    validated_ast.children[0].should.be.a(UpdateExpressionAddAction)
+
+
+def test_normalize_with_multiple_actions__order_is_preserved(table):
+    update_expression = "ADD s :value REMOVE a[3], a[1], a[2] SET t=:value"
+    update_expression_ast = UpdateExpressionParser.make(update_expression)
+    item = Item(
+        hash_key=DynamoType({"S": "id"}),
+        range_key=None,
+        attrs={
+            "id": {"S": "foo2"},
+            "a": {"L": [{"S": "val1"}, {"S": "val2"}, {"S": "val3"}, {"S": "val4"}]},
+            "s": {"SS": ["value1", "value2", "value3"]},
+        },
+    )
+    validated_ast = UpdateExpressionValidator(
+        update_expression_ast,
+        expression_attribute_names=None,
+        expression_attribute_values={":value": {"SS": ["value2", "value5"]}},
+        item=item,
+        table=table,
+    ).validate()
+    validated_ast.children.should.have.length_of(2)
+    # add clause first
+    validated_ast.children[0].should.be.a(UpdateExpressionAddClause)
+    # rest of the expression next
+    validated_ast.children[1].should.be.a(UpdateExpression)
+
+    validated_ast.normalize()
+    validated_ast.children.should.have.length_of(5)
+    # add action first
+    validated_ast.children[0].should.be.a(UpdateExpressionAddAction)
+    # Removal actions in reverse order
+    validated_ast.children[1].should.be.a(UpdateExpressionRemoveAction)
+    validated_ast.children[1]._get_value().should.equal(3)
+    validated_ast.children[2].should.be.a(UpdateExpressionRemoveAction)
+    validated_ast.children[2]._get_value().should.equal(2)
+    validated_ast.children[3].should.be.a(UpdateExpressionRemoveAction)
+    validated_ast.children[3]._get_value().should.equal(1)
+    # Set action last, as per insertion order
+    validated_ast.children[4].should.be.a(UpdateExpressionSetAction)