From 58df83f39fd8c09da4a1cc008026f4b2559c0a54 Mon Sep 17 00:00:00 2001
From: Bert Blommers <bblommers@users.noreply.github.com>
Date: Tue, 12 Oct 2021 19:32:10 +0000
Subject: [PATCH] DynamoDB - Validate the nr of Add-clauses on update (#4398)

---
 moto/dynamodb2/exceptions.py                  |  7 ++++++
 moto/dynamodb2/models/__init__.py             |  1 +
 moto/dynamodb2/parsing/ast_nodes.py           | 16 ++++++++++++++
 .../test_dynamodb_exceptions.py               | 22 +++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/moto/dynamodb2/exceptions.py b/moto/dynamodb2/exceptions.py
index 7d517d8c1..40c5591a5 100644
--- a/moto/dynamodb2/exceptions.py
+++ b/moto/dynamodb2/exceptions.py
@@ -204,3 +204,10 @@ class UpdateHashRangeKeyException(MockValidationException):
 
     def __init__(self, key_name):
         super(UpdateHashRangeKeyException, self).__init__(self.msg.format(key_name))
+
+
+class TooManyAddClauses(InvalidUpdateExpression):
+    msg = 'The "ADD" section can only be used once in an update expression;'
+
+    def __init__(self):
+        super(TooManyAddClauses, self).__init__(self.msg)
diff --git a/moto/dynamodb2/models/__init__.py b/moto/dynamodb2/models/__init__.py
index 4cff16329..3bd248953 100644
--- a/moto/dynamodb2/models/__init__.py
+++ b/moto/dynamodb2/models/__init__.py
@@ -1388,6 +1388,7 @@ class DynamoDBBackend(BaseBackend):
             # Parse expression to get validation errors
             update_expression_ast = UpdateExpressionParser.make(update_expression)
             update_expression = re.sub(r"\s*([=\+-])\s*", "\\1", update_expression)
+            update_expression_ast.validate()
 
         if all([table.hash_key_attr in key, table.range_key_attr in key]):
             # Covers cases where table has hash and range keys, ``key`` param
diff --git a/moto/dynamodb2/parsing/ast_nodes.py b/moto/dynamodb2/parsing/ast_nodes.py
index 4b1907dad..66bfeb4f6 100644
--- a/moto/dynamodb2/parsing/ast_nodes.py
+++ b/moto/dynamodb2/parsing/ast_nodes.py
@@ -3,6 +3,7 @@ from abc import abstractmethod
 from collections import deque
 
 from moto.dynamodb2.models import DynamoType
+from ..exceptions import TooManyAddClauses
 
 
 class Node(metaclass=abc.ABCMeta):
@@ -20,6 +21,21 @@ class Node(metaclass=abc.ABCMeta):
     def set_parent(self, parent_node):
         self.parent = parent_node
 
+    def validate(self):
+        if self.type == "UpdateExpression":
+            nr_of_clauses = len(self.find_clauses(UpdateExpressionAddClause))
+            if nr_of_clauses > 1:
+                raise TooManyAddClauses()
+
+    def find_clauses(self, clause_type):
+        clauses = []
+        for child in self.children or []:
+            if isinstance(child, clause_type):
+                clauses.append(child)
+            elif isinstance(child, Expression):
+                clauses.extend(child.find_clauses(clause_type))
+        return clauses
+
 
 class LeafNode(Node):
     """A LeafNode is a Node where none of the children are Nodes themselves."""
diff --git a/tests/test_dynamodb2/test_dynamodb_exceptions.py b/tests/test_dynamodb2/test_dynamodb_exceptions.py
index c2f8a842d..377ef2e0e 100644
--- a/tests/test_dynamodb2/test_dynamodb_exceptions.py
+++ b/tests/test_dynamodb2/test_dynamodb_exceptions.py
@@ -147,3 +147,25 @@ def test_empty_expressionattributenames_with_projection():
     err = exc.value.response["Error"]
     err["Code"].should.equal("ValidationException")
     err["Message"].should.equal("ExpressionAttributeNames must not be empty")
+
+
+@mock_dynamodb2
+def test_update_item_range_key_set():
+    ddb = boto3.resource("dynamodb", region_name="us-east-1")
+
+    # Create the DynamoDB table.
+    table = ddb.create_table(
+        TableName="test-table", BillingMode="PAY_PER_REQUEST", **table_schema
+    )
+
+    with pytest.raises(ClientError) as exc:
+        table.update_item(
+            Key={"partitionKey": "the-key"},
+            UpdateExpression="ADD x :one SET a = :a ADD y :one",
+            ExpressionAttributeValues={":one": 1, ":a": "lore ipsum"},
+        )
+    err = exc.value.response["Error"]
+    err["Code"].should.equal("ValidationException")
+    err["Message"].should.equal(
+        'Invalid UpdateExpression: The "ADD" section can only be used once in an update expression;'
+    )