From cc568c16565e5e9e895a73ce6389b58b5406b0a5 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Sun, 29 Aug 2021 19:04:42 +0100 Subject: [PATCH] DynamoDB - Improve handling of ReturnConsumedCapacity-param (#4241) --- moto/dynamodb2/responses.py | 60 ++++++- tests/test_dynamodb2/test_dynamodb.py | 157 ------------------ .../test_dynamodb_consumedcapacity.py | 147 ++++++++++++++++ 3 files changed, 202 insertions(+), 162 deletions(-) create mode 100644 tests/test_dynamodb2/test_dynamodb_consumedcapacity.py diff --git a/moto/dynamodb2/responses.py b/moto/dynamodb2/responses.py index 288ff27c1..486c000e5 100644 --- a/moto/dynamodb2/responses.py +++ b/moto/dynamodb2/responses.py @@ -5,6 +5,7 @@ import json import re import itertools +from functools import wraps from moto.core.responses import BaseResponse from moto.core.utils import camelcase_to_underscores, amz_crc32, amzn_request_id @@ -19,6 +20,55 @@ from moto.dynamodb2.models import dynamodb_backends, dynamo_json_dump TRANSACTION_MAX_ITEMS = 25 +def include_consumed_capacity(val=1.0): + def _inner(f): + @wraps(f) + def _wrapper(*args, **kwargs): + (handler,) = args + expected_capacity = handler.body.get("ReturnConsumedCapacity", "NONE") + if expected_capacity not in ["NONE", "TOTAL", "INDEXES"]: + type_ = "ValidationException" + message = "1 validation error detected: Value '{}' at 'returnConsumedCapacity' failed to satisfy constraint: Member must satisfy enum value set: [INDEXES, TOTAL, NONE]".format( + expected_capacity + ) + return ( + 400, + handler.response_headers, + dynamo_json_dump({"__type": type_, "message": message}), + ) + table_name = handler.body.get("TableName", "") + index_name = handler.body.get("IndexName", None) + + response = f(*args, **kwargs) + + if isinstance(response, str): + body = json.loads(response) + + if expected_capacity == "TOTAL": + body["ConsumedCapacity"] = { + "TableName": table_name, + "CapacityUnits": val, + } + elif expected_capacity == "INDEXES": + body["ConsumedCapacity"] = { + "TableName": table_name, + "CapacityUnits": val, + "Table": {"CapacityUnits": val}, + } + if index_name: + body["ConsumedCapacity"]["LocalSecondaryIndexes"] = { + index_name: {"CapacityUnits": val} + } + + return dynamo_json_dump(body) + + return response + + return _wrapper + + return _inner + + def put_has_empty_keys(field_updates, table): if table: key_names = table.key_attributes @@ -254,6 +304,7 @@ class DynamoHandler(BaseResponse): er = "com.amazonaws.dynamodb.v20111205#ResourceNotFoundException" return self.error(er, "Requested resource not found") + @include_consumed_capacity() def put_item(self): name = self.body["TableName"] item = self.body["Item"] @@ -310,7 +361,6 @@ class DynamoHandler(BaseResponse): if result: item_dict = result.to_json() - item_dict["ConsumedCapacity"] = {"TableName": name, "CapacityUnits": 1} if return_values == "ALL_OLD": item_dict["Attributes"] = existing_attributes else: @@ -349,6 +399,7 @@ class DynamoHandler(BaseResponse): return dynamo_json_dump(response) + @include_consumed_capacity(0.5) def get_item(self): name = self.body["TableName"] table = self.dynamodb_backend.get_table(name) @@ -372,11 +423,10 @@ class DynamoHandler(BaseResponse): return self.error(er, "Validation Exception") if item: item_dict = item.describe_attrs(attributes=None) - item_dict["ConsumedCapacity"] = {"TableName": name, "CapacityUnits": 0.5} return dynamo_json_dump(item_dict) else: # Item not found - return 200, self.response_headers, "{}" + return dynamo_json_dump({}) def batch_get_item(self): table_batches = self.body["RequestItems"] @@ -441,6 +491,7 @@ class DynamoHandler(BaseResponse): unique_keys.append(k) return False + @include_consumed_capacity() def query(self): name = self.body["TableName"] key_condition_expression = self.body.get("KeyConditionExpression") @@ -618,7 +669,6 @@ class DynamoHandler(BaseResponse): result = { "Count": len(items), - "ConsumedCapacity": {"TableName": name, "CapacityUnits": 1}, "ScannedCount": scanned_count, } @@ -649,6 +699,7 @@ class DynamoHandler(BaseResponse): return projection_expression + @include_consumed_capacity() def scan(self): name = self.body["TableName"] @@ -700,7 +751,6 @@ class DynamoHandler(BaseResponse): result = { "Count": len(items), "Items": [item.attrs for item in items], - "ConsumedCapacity": {"TableName": name, "CapacityUnits": 1}, "ScannedCount": scanned_count, } if last_evaluated_key is not None: diff --git a/tests/test_dynamodb2/test_dynamodb.py b/tests/test_dynamodb2/test_dynamodb.py index 01984df87..6fe989475 100644 --- a/tests/test_dynamodb2/test_dynamodb.py +++ b/tests/test_dynamodb2/test_dynamodb.py @@ -589,42 +589,6 @@ def test_query_invalid_table(): assert exception.response["Error"]["Code"] == "ResourceNotFoundException" -@requires_boto_gte("2.9") -@mock_dynamodb2 -def test_scan_returns_consumed_capacity(): - name = "TestTable" - conn = boto3.client( - "dynamodb", - region_name="us-west-2", - aws_access_key_id="ak", - aws_secret_access_key="sk", - ) - - conn.create_table( - TableName=name, - KeySchema=[{"AttributeName": "forum_name", "KeyType": "HASH"}], - AttributeDefinitions=[{"AttributeName": "forum_name", "AttributeType": "S"}], - ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, - ) - - conn.put_item( - TableName=name, - Item={ - "forum_name": {"S": "LOLCat Forum"}, - "subject": {"S": "Check this out!"}, - "Body": {"S": "http://url_to_lolcat.gif"}, - "SentBy": {"S": "test"}, - "ReceivedTime": {"S": "12/9/2011 11:36:03 PM"}, - }, - ) - - response = conn.scan(TableName=name) - - assert "ConsumedCapacity" in response - assert "CapacityUnits" in response["ConsumedCapacity"] - assert response["ConsumedCapacity"]["TableName"] == name - - @requires_boto_gte("2.9") @mock_dynamodb2 def test_put_item_with_special_chars(): @@ -710,37 +674,6 @@ def test_put_item_with_streams(): stream_record["dynamodb"]["SizeBytes"].should.be.equal(447) -@requires_boto_gte("2.9") -@mock_dynamodb2 -def test_query_returns_consumed_capacity(): - dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - - # Create the DynamoDB table. - table = dynamodb.create_table( - TableName="users", - KeySchema=[ - {"AttributeName": "forum_name", "KeyType": "HASH"}, - {"AttributeName": "subject", "KeyType": "RANGE"}, - ], - AttributeDefinitions=[ - {"AttributeName": "forum_name", "AttributeType": "S"}, - {"AttributeName": "subject", "AttributeType": "S"}, - ], - ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, - ) - table = dynamodb.Table("users") - - table.put_item( - Item={"forum_name": "the-key", "subject": "123", "body": "some test message"} - ) - - results = table.query(KeyConditionExpression=Key("forum_name").eq("the-key")) - - assert "ConsumedCapacity" in results - assert "CapacityUnits" in results["ConsumedCapacity"] - assert results["ConsumedCapacity"]["CapacityUnits"] == 1 - - @mock_dynamodb2 def test_basic_projection_expression_using_get_item(): dynamodb = boto3.resource("dynamodb", region_name="us-east-1") @@ -1435,96 +1368,6 @@ def test_nested_projection_expression_using_scan_with_attr_expression_names(): ) -@mock_dynamodb2 -def test_put_item_returns_consumed_capacity(): - dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - - # Create the DynamoDB table. - table = dynamodb.create_table( - TableName="users", - KeySchema=[ - {"AttributeName": "forum_name", "KeyType": "HASH"}, - {"AttributeName": "subject", "KeyType": "RANGE"}, - ], - AttributeDefinitions=[ - {"AttributeName": "forum_name", "AttributeType": "S"}, - {"AttributeName": "subject", "AttributeType": "S"}, - ], - ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, - ) - table = dynamodb.Table("users") - - response = table.put_item( - Item={"forum_name": "the-key", "subject": "123", "body": "some test message"} - ) - - assert "ConsumedCapacity" in response - - -@mock_dynamodb2 -def test_update_item_returns_consumed_capacity(): - dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - - # Create the DynamoDB table. - table = dynamodb.create_table( - TableName="users", - KeySchema=[ - {"AttributeName": "forum_name", "KeyType": "HASH"}, - {"AttributeName": "subject", "KeyType": "RANGE"}, - ], - AttributeDefinitions=[ - {"AttributeName": "forum_name", "AttributeType": "S"}, - {"AttributeName": "subject", "AttributeType": "S"}, - ], - ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, - ) - table = dynamodb.Table("users") - - table.put_item( - Item={"forum_name": "the-key", "subject": "123", "body": "some test message"} - ) - - response = table.update_item( - Key={"forum_name": "the-key", "subject": "123"}, - UpdateExpression="set body=:tb", - ExpressionAttributeValues={":tb": "a new message"}, - ) - - assert "ConsumedCapacity" in response - assert "CapacityUnits" in response["ConsumedCapacity"] - assert "TableName" in response["ConsumedCapacity"] - - -@mock_dynamodb2 -def test_get_item_returns_consumed_capacity(): - dynamodb = boto3.resource("dynamodb", region_name="us-east-1") - - # Create the DynamoDB table. - table = dynamodb.create_table( - TableName="users", - KeySchema=[ - {"AttributeName": "forum_name", "KeyType": "HASH"}, - {"AttributeName": "subject", "KeyType": "RANGE"}, - ], - AttributeDefinitions=[ - {"AttributeName": "forum_name", "AttributeType": "S"}, - {"AttributeName": "subject", "AttributeType": "S"}, - ], - ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, - ) - table = dynamodb.Table("users") - - table.put_item( - Item={"forum_name": "the-key", "subject": "123", "body": "some test message"} - ) - - response = table.get_item(Key={"forum_name": "the-key", "subject": "123"}) - - assert "ConsumedCapacity" in response - assert "CapacityUnits" in response["ConsumedCapacity"] - assert "TableName" in response["ConsumedCapacity"] - - @mock_dynamodb2 def test_put_empty_item(): dynamodb = boto3.resource("dynamodb", region_name="us-east-1") diff --git a/tests/test_dynamodb2/test_dynamodb_consumedcapacity.py b/tests/test_dynamodb2/test_dynamodb_consumedcapacity.py new file mode 100644 index 000000000..03e8d0249 --- /dev/null +++ b/tests/test_dynamodb2/test_dynamodb_consumedcapacity.py @@ -0,0 +1,147 @@ +import boto3 +import pytest +import sure # noqa + +from botocore.exceptions import ClientError +from moto import mock_dynamodb2 + + +@mock_dynamodb2 +def test_error_on_wrong_value_for_consumed_capacity(): + resource = boto3.resource("dynamodb", region_name="ap-northeast-3") + client = boto3.client("dynamodb", region_name="ap-northeast-3") + client.create_table( + TableName="jobs", + KeySchema=[{"AttributeName": "job_id", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "job_id", "AttributeType": "S"}], + ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, + ) + + table = resource.Table("jobs") + item = {"job_id": "asdasdasd", "expires_at": "1"} + + # PUT_ITEM + with pytest.raises(ClientError) as ex: + table.put_item(Item=item, ReturnConsumedCapacity="Garbage") + err = ex.value.response["Error"] + err["Code"].should.equal("ValidationException") + err["Message"].should.equal( + "1 validation error detected: Value 'Garbage' at 'returnConsumedCapacity' failed to satisfy constraint: Member must satisfy enum value set: [INDEXES, TOTAL, NONE]" + ) + + +@mock_dynamodb2 +def test_consumed_capacity_get_unknown_item(): + conn = boto3.client("dynamodb", region_name="us-east-1") + conn.create_table( + TableName="test_table", + KeySchema=[{"AttributeName": "u", "KeyType": "HASH"}], + AttributeDefinitions=[{"AttributeName": "u", "AttributeType": "S"}], + ) + response = conn.get_item( + TableName="test_table", + Key={"u": {"S": "does_not_exist"}}, + ReturnConsumedCapacity="TOTAL", + ) + + # Should still return ConsumedCapacity, even if it does not return an item + response.should.have.key("ConsumedCapacity") + response["ConsumedCapacity"].should.equal( + {"TableName": "test_table", "CapacityUnits": 0.5} + ) + + +@mock_dynamodb2 +@pytest.mark.parametrize( + "capacity,should_have_capacity,should_have_table", + [ + [None, False, False], + ["NONE", False, False], + ["TOTAL", True, False], + ["INDEXES", True, True], + ], +) +def test_only_return_consumed_capacity_when_required( + capacity, should_have_capacity, should_have_table +): + resource = boto3.resource("dynamodb", region_name="ap-northeast-3") + client = boto3.client("dynamodb", region_name="ap-northeast-3") + client.create_table( + TableName="jobs", + KeySchema=[{"AttributeName": "job_id", "KeyType": "HASH"}], + LocalSecondaryIndexes=[ + { + "IndexName": "job_name-index", + "KeySchema": [{"AttributeName": "job_name", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + } + ], + AttributeDefinitions=[ + {"AttributeName": "job_id", "AttributeType": "S"}, + {"AttributeName": "job_name", "AttributeType": "S"}, + ], + ProvisionedThroughput={"ReadCapacityUnits": 5, "WriteCapacityUnits": 5}, + ) + + table = resource.Table("jobs") + item = {"job_id": "asdasdasd", "expires_at": "1"} + + # PUT_ITEM + args = {"Item": item} + if capacity: + args["ReturnConsumedCapacity"] = capacity + response = table.put_item(**args) + validate_response(response, should_have_capacity, should_have_table) + + # GET_ITEM + args = {"Key": item} + if capacity: + args["ReturnConsumedCapacity"] = capacity + response = table.get_item(**args) + validate_response(response, should_have_capacity, should_have_table, value=0.5) + + # SCAN + args = {"TableName": "jobs"} + if capacity: + args["ReturnConsumedCapacity"] = capacity + response = client.scan(**args) + validate_response(response, should_have_capacity, should_have_table) + + # SCAN_INDEX + args["IndexName"] = "job_name-index" + response = client.scan(**args) + validate_response(response, should_have_capacity, should_have_table, is_index=True) + + # QUERY + args = { + "TableName": "jobs", + "KeyConditionExpression": "job_id = :id", + "ExpressionAttributeValues": {":id": {"S": "asdasdasd"}}, + } + if capacity: + args["ReturnConsumedCapacity"] = capacity + response = client.query(**args) + validate_response(response, should_have_capacity, should_have_table) + + # QUERY_INDEX + args["IndexName"] = "job_name-index" + response = client.query(**args) + validate_response(response, should_have_capacity, should_have_table, is_index=True) + + +def validate_response( + response, should_have_capacity, should_have_table, is_index=False, value=1.0 +): + if should_have_capacity: + response.should.have.key("ConsumedCapacity") + response["ConsumedCapacity"]["TableName"].should.equal("jobs") + response["ConsumedCapacity"]["CapacityUnits"].should.equal(value) + if should_have_table: + response["ConsumedCapacity"]["Table"].should.equal({"CapacityUnits": value}) + if is_index: + response["ConsumedCapacity"].should.have.key("LocalSecondaryIndexes") + response["ConsumedCapacity"]["LocalSecondaryIndexes"].should.equal( + {"job_name-index": {"CapacityUnits": value}} + ) + else: + response.shouldnt.have.key("ConsumedCapacity")