DynamoDB: Improve parsing of UpdateExpressions with multiple REMOVE c…

…lauses (#8120)

moto/dynamodb/exceptions.py

@@ -305,11 +305,11 @@ def __init__(self, names: List[str]):
         )
 
 
-class TooManyAddClauses(InvalidUpdateExpression):
-    msg = 'The "ADD" section can only be used once in an update expression;'
-
-    def __init__(self) -> None:
-        super().__init__(self.msg)
+class TooManyClauses(InvalidUpdateExpression):
+    def __init__(self, _type: str) -> None:
+        super().__init__(
+            f'The "{_type}" section can only be used once in an update expression;'
+        )
 
 
 class ResourceNotFoundException(JsonRESTError):

moto/dynamodb/parsing/ast_nodes.py

@@ -8,7 +8,7 @@
 from ..exceptions import (
     DuplicateUpdateExpression,
     MockValidationException,
-    TooManyAddClauses,
+    TooManyClauses,
 )
 
 
@@ -29,9 +29,10 @@ def set_parent(self, parent_node):
 
     def validate(self, limit_set_actions: bool = False) -> None:
         if self.type == "UpdateExpression":
-            nr_of_clauses = len(self.find_clauses([UpdateExpressionAddClause]))
-            if nr_of_clauses > 1:
-                raise TooManyAddClauses()
+            if len(self.find_clauses([UpdateExpressionAddClause])) > 1:
+                raise TooManyClauses("ADD")
+            if len(self.find_clauses([UpdateExpressionRemoveClause])) > 1:
+                raise TooManyClauses("REMOVE")
             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]
@@ -176,8 +177,11 @@ def _get_value(self):
 
     def __lt__(self, other):
         self_value = self._get_value()
-
-        return self_value < other._get_value()
+        other_value = other._get_value()
+        if isinstance(self_value, int) and isinstance(other_value, int):
+            return self_value < other_value
+        else:
+            return str(self_value) < str(other_value)
 
 
 class UpdateExpressionAddActions(UpdateExpressionClause):

tests/test_dynamodb/test_dynamodb_update_expressions.py

@@ -225,3 +225,65 @@ def test_update_item_with_empty_expression(table_name=None):
     assert (
         err["Message"] == "Invalid UpdateExpression: The expression can not be empty;"
     )
+
+
+@pytest.mark.aws_verified
+@dynamodb_aws_verified()
+def test_update_expression_with_multiple_remove_clauses(table_name=None):
+    ddb_client = boto3.client("dynamodb", "us-east-1")
+    payload = {
+        "pk": {"S": "primary_key"},
+        "current_user": {"M": {"name": {"S": "John"}, "surname": {"S": "Doe"}}},
+        "user_list": {
+            "L": [
+                {"M": {"name": {"S": "John"}, "surname": {"S": "Doe"}}},
+                {"M": {"name": {"S": "Jane"}, "surname": {"S": "Smith"}}},
+            ]
+        },
+        "some_param": {"NULL": True},
+    }
+    ddb_client.put_item(TableName=table_name, Item=payload)
+    with pytest.raises(ClientError) as exc:
+        ddb_client.update_item(
+            TableName=table_name,
+            Key={"pk": {"S": "primary_key"}},
+            UpdateExpression="REMOVE #ulist[0] SET current_user = :current_user REMOVE some_param",
+            ExpressionAttributeNames={"#ulist": "user_list"},
+            ExpressionAttributeValues={":current_user": {"M": {"name": {"S": "Jane"}}}},
+        )
+    err = exc.value.response["Error"]
+    assert err["Code"] == "ValidationException"
+    assert (
+        err["Message"]
+        == 'Invalid UpdateExpression: The "REMOVE" section can only be used once in an update expression;'
+    )
+
+
+@pytest.mark.aws_verified
+@dynamodb_aws_verified()
+def test_update_expression_remove_list_and_attribute(table_name=None):
+    ddb_client = boto3.client("dynamodb", "us-east-1")
+    payload = {
+        "pk": {"S": "primary_key"},
+        "user_list": {
+            "L": [
+                {"M": {"name": {"S": "John"}, "surname": {"S": "Doe"}}},
+                {"M": {"name": {"S": "Jane"}, "surname": {"S": "Smith"}}},
+            ]
+        },
+        "some_param": {"NULL": True},
+    }
+    ddb_client.put_item(TableName=table_name, Item=payload)
+    ddb_client.update_item(
+        TableName=table_name,
+        Key={"pk": {"S": "primary_key"}},
+        UpdateExpression="REMOVE #ulist[0], some_param",
+        ExpressionAttributeNames={"#ulist": "user_list"},
+    )
+    item = ddb_client.get_item(
+        TableName=table_name, Key={"pk": {"S": "primary_key"}, "sk": {"S": "sort_key"}}
+    )["Item"]
+    assert item == {
+        "user_list": {"L": [{"M": {"name": {"S": "Jane"}, "surname": {"S": "Smith"}}}]},
+        "pk": {"S": "primary_key"},
+    }