From cb27b55008ca90f048e3d14612621971385d6e26 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Wed, 14 Dec 2022 10:07:34 -0100 Subject: [PATCH] DynamoDB: Allow removal of multiple listitems (#5767) --- moto/dynamodb/models/__init__.py | 1 + moto/dynamodb/parsing/ast_nodes.py | 60 ++++++++++++++-- moto/dynamodb/parsing/executors.py | 2 + tests/test_dynamodb/test_dynamodb.py | 23 ++++++ tests/test_dynamodb/test_dynamodb_executor.py | 70 +++++++++++++++++++ 5 files changed, 151 insertions(+), 5 deletions(-) diff --git a/moto/dynamodb/models/__init__.py b/moto/dynamodb/models/__init__.py index d02a2c5c6..cd06d539b 100644 --- a/moto/dynamodb/models/__init__.py +++ b/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 diff --git a/moto/dynamodb/parsing/ast_nodes.py b/moto/dynamodb/parsing/ast_nodes.py index 4a2ba1d36..a021d3ca6 100644 --- a/moto/dynamodb/parsing/ast_nodes.py +++ b/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): """ diff --git a/moto/dynamodb/parsing/executors.py b/moto/dynamodb/parsing/executors.py index 2c0be68b9..9b8df9602 100644 --- a/moto/dynamodb/parsing/executors.py +++ b/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): diff --git a/tests/test_dynamodb/test_dynamodb.py b/tests/test_dynamodb/test_dynamodb.py index 873c8738e..244502cdb 100644 --- a/tests/test_dynamodb/test_dynamodb.py +++ b/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" diff --git a/tests/test_dynamodb/test_dynamodb_executor.py b/tests/test_dynamodb/test_dynamodb_executor.py index 5578a5bb3..75c5d84c0 100644 --- a/tests/test_dynamodb/test_dynamodb_executor.py +++ b/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)