From 9e7803dc3601427b8f195f715f07a76a23216749 Mon Sep 17 00:00:00 2001 From: pvbouwel Date: Thu, 7 May 2020 21:29:20 +0100 Subject: [PATCH] [Bugfix] UpdateExpression using ADD from zero (#2975) When using the ADD syntax to sum up different components the path that is provided is allowed to be non-existent. In such a case DynamoDB will initialize it depending on the type of the value. If it is a number it will be initialized with 0. If it is a set it will be initialized with an empty set. --- moto/dynamodb2/parsing/executors.py | 25 +++++- tests/test_dynamodb2/test_dynamodb.py | 78 +++++++++++++++++++ .../test_dynamodb_table_with_range_key.py | 23 +++--- 3 files changed, 112 insertions(+), 14 deletions(-) diff --git a/moto/dynamodb2/parsing/executors.py b/moto/dynamodb2/parsing/executors.py index 8c51c9cec..2f2f2bb82 100644 --- a/moto/dynamodb2/parsing/executors.py +++ b/moto/dynamodb2/parsing/executors.py @@ -1,6 +1,10 @@ from abc import abstractmethod -from moto.dynamodb2.exceptions import IncorrectOperandType, IncorrectDataType +from moto.dynamodb2.exceptions import ( + IncorrectOperandType, + IncorrectDataType, + ProvidedKeyDoesNotExist, +) from moto.dynamodb2.models import DynamoType from moto.dynamodb2.models.dynamo_type import DDBTypeConversion, DDBType from moto.dynamodb2.parsing.ast_nodes import ( @@ -193,7 +197,18 @@ class AddExecutor(NodeExecutor): value_to_add = self.get_action_value() if isinstance(value_to_add, DynamoType): if value_to_add.is_set(): - current_string_set = self.get_item_at_end_of_path(item) + try: + current_string_set = self.get_item_at_end_of_path(item) + except ProvidedKeyDoesNotExist: + current_string_set = DynamoType({value_to_add.type: []}) + SetExecutor.set( + item_part_to_modify_with_set=self.get_item_before_end_of_path( + item + ), + element_to_set=self.get_element_to_action(), + value_to_set=current_string_set, + expression_attribute_names=self.expression_attribute_names, + ) assert isinstance(current_string_set, DynamoType) if not current_string_set.type == value_to_add.type: raise IncorrectDataType() @@ -204,7 +219,11 @@ class AddExecutor(NodeExecutor): else: current_string_set.value.append(value) elif value_to_add.type == DDBType.NUMBER: - existing_value = self.get_item_at_end_of_path(item) + try: + existing_value = self.get_item_at_end_of_path(item) + except ProvidedKeyDoesNotExist: + existing_value = DynamoType({DDBType.NUMBER: "0"}) + assert isinstance(existing_value, DynamoType) if not existing_value.type == DDBType.NUMBER: raise IncorrectDataType() diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index 470c5f8ff..9f917a7ae 100644 --- a/tests/test_dynamodb2/test_dynamodb.py +++ b/tests/test_dynamodb2/test_dynamodb.py @@ -5029,3 +5029,81 @@ def test_update_item_atomic_counter_return_values(): "v" in response["Attributes"] ), "v has been updated, and should be returned here" response["Attributes"]["v"]["N"].should.equal("8") + + +@mock_dynamodb2 +def test_update_item_atomic_counter_from_zero(): + table = "table_t" + ddb_mock = boto3.client("dynamodb", region_name="eu-west-1") + ddb_mock.create_table( + TableName=table, + KeySchema=[{"AttributeName": "t_id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "t_id", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + + key = {"t_id": {"S": "item1"}} + + ddb_mock.put_item( + TableName=table, Item=key, + ) + + ddb_mock.update_item( + TableName=table, + Key=key, + UpdateExpression="add n_i :inc1, n_f :inc2", + ExpressionAttributeValues={":inc1": {"N": "1.2"}, ":inc2": {"N": "-0.5"}}, + ) + updated_item = ddb_mock.get_item(TableName=table, Key=key)["Item"] + assert updated_item["n_i"]["N"] == "1.2" + assert updated_item["n_f"]["N"] == "-0.5" + + +@mock_dynamodb2 +def test_update_item_add_to_non_existent_set(): + table = "table_t" + ddb_mock = boto3.client("dynamodb", region_name="eu-west-1") + ddb_mock.create_table( + TableName=table, + KeySchema=[{"AttributeName": "t_id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "t_id", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + key = {"t_id": {"S": "item1"}} + ddb_mock.put_item( + TableName=table, Item=key, + ) + + ddb_mock.update_item( + TableName=table, + Key=key, + UpdateExpression="add s_i :s1", + ExpressionAttributeValues={":s1": {"SS": ["hello"]}}, + ) + updated_item = ddb_mock.get_item(TableName=table, Key=key)["Item"] + assert updated_item["s_i"]["SS"] == ["hello"] + + +@mock_dynamodb2 +def test_update_item_add_to_non_existent_number_set(): + table = "table_t" + ddb_mock = boto3.client("dynamodb", region_name="eu-west-1") + ddb_mock.create_table( + TableName=table, + KeySchema=[{"AttributeName": "t_id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "t_id", "AttributeType": "S"}], + BillingMode="PAY_PER_REQUEST", + ) + key = {"t_id": {"S": "item1"}} + ddb_mock.put_item( + TableName=table, Item=key, + ) + + ddb_mock.update_item( + TableName=table, + Key=key, + UpdateExpression="add s_i :s1", + ExpressionAttributeValues={":s1": {"NS": ["3"]}}, + ) + updated_item = ddb_mock.get_item(TableName=table, Key=key)["Item"] + assert updated_item["s_i"]["NS"] == ["3"] diff --git a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py index 6fba713ec..33f65d5ec 100644 --- a/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py +++ b/tests/test_dynamodb2/test_dynamodb_table_with_range_key.py @@ -1307,16 +1307,16 @@ def test_update_item_add_with_expression(): ExpressionAttributeValues={":v": {"item4"}}, ) current_item["str_set"] = current_item["str_set"].union({"item4"}) - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Update item to add a string value to a non-existing set - # Should throw: 'The provided key element does not match the schema' - assert_failure_due_to_key_not_in_schema( - table.update_item, + table.update_item( Key=item_key, UpdateExpression="ADD non_existing_str_set :v", ExpressionAttributeValues={":v": {"item4"}}, ) + current_item["non_existing_str_set"] = {"item4"} + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Update item to add a num value to a num set table.update_item( @@ -1325,7 +1325,7 @@ def test_update_item_add_with_expression(): ExpressionAttributeValues={":v": {6}}, ) current_item["num_set"] = current_item["num_set"].union({6}) - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Update item to add a value to a number value table.update_item( @@ -1334,7 +1334,7 @@ def test_update_item_add_with_expression(): ExpressionAttributeValues={":v": 20}, ) current_item["num_val"] = current_item["num_val"] + 20 - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Attempt to add a number value to a string set, should raise Client Error table.update_item.when.called_with( @@ -1342,7 +1342,7 @@ def test_update_item_add_with_expression(): UpdateExpression="ADD str_set :v", ExpressionAttributeValues={":v": 20}, ).should.have.raised(ClientError) - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Attempt to add a number set to the string set, should raise a ClientError table.update_item.when.called_with( @@ -1350,7 +1350,7 @@ def test_update_item_add_with_expression(): UpdateExpression="ADD str_set :v", ExpressionAttributeValues={":v": {20}}, ).should.have.raised(ClientError) - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Attempt to update with a bad expression table.update_item.when.called_with( @@ -1388,17 +1388,18 @@ def test_update_item_add_with_nested_sets(): current_item["nested"]["str_set"] = current_item["nested"]["str_set"].union( {"item4"} ) - dict(table.get_item(Key=item_key)["Item"]).should.equal(current_item) + assert dict(table.get_item(Key=item_key)["Item"]) == current_item # Update item to add a string value to a non-existing set # Should raise - assert_failure_due_to_key_not_in_schema( - table.update_item, + table.update_item( Key=item_key, UpdateExpression="ADD #ns.#ne :v", ExpressionAttributeNames={"#ns": "nested", "#ne": "non_existing_str_set"}, ExpressionAttributeValues={":v": {"new_item"}}, ) + current_item["nested"]["non_existing_str_set"] = {"new_item"} + assert dict(table.get_item(Key=item_key)["Item"]) == current_item @mock_dynamodb2