#1874 - Count item size based on contents of actual dictionary

moto/dynamodb2/exceptions.py

@@ -1,2 +1,7 @@
 class InvalidIndexNameError(ValueError):
     pass
+
+
+class ItemSizeTooLarge(Exception):
+    message = 'Item size has exceeded the maximum allowed size'
+    pass

moto/dynamodb2/models.py

@@ -16,7 +16,7 @@ from moto.core.exceptions import JsonRESTError
 from .comparisons import get_comparison_func
 from .comparisons import get_filter_expression
 from .comparisons import get_expected
-from .exceptions import InvalidIndexNameError
+from .exceptions import InvalidIndexNameError, ItemSizeTooLarge
 
 
 class DynamoJsonEncoder(json.JSONEncoder):
@@ -30,6 +30,10 @@ def dynamo_json_dump(dynamo_object):
     return json.dumps(dynamo_object, cls=DynamoJsonEncoder)
 
 
+def bytesize(val):
+    return len(str(val).encode('utf-8'))
+
+
 class DynamoType(object):
     """
     http://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DataModel.html#DataModelDataTypes
@@ -99,6 +103,22 @@ class DynamoType(object):
 
         return None
 
+    def size(self):
+        if self.is_number():
+            value_size = len(str(self.value))
+        elif self.is_set():
+            sub_type = self.type[0]
+            value_size = sum([DynamoType({sub_type: v}).size() for v in self.value])
+        elif self.is_list():
+            value_size = sum([DynamoType(v).size() for v in self.value])
+        elif self.is_map():
+            value_size = sum([bytesize(k) + DynamoType(v).size() for k, v in self.value.items()])
+        elif type(self.value) == bool:
+            value_size = 1
+        else:
+            value_size = bytesize(self.value)
+        return value_size
+
     def to_json(self):
         return {self.type: self.value}
 
@@ -126,6 +146,39 @@ class DynamoType(object):
         return self.type == other.type
 
 
+# https://github.com/spulec/moto/issues/1874
+# Ensure that the total size of an item does not exceed 400kb
+class LimitedSizeDict(dict):
+    def __init__(self, *args, **kwargs):
+        self.update(*args, **kwargs)
+
+    def __setitem__(self, key, value):
+        current_item_size = sum([item.size() if type(item) == DynamoType else bytesize(str(item)) for item in (self.keys() + self.values())])
+        new_item_size = bytesize(key) + (value.size() if type(value) == DynamoType else bytesize(str(value)))
+        # Official limit is set to 400000 (400KB)
+        # Manual testing confirms that the actual limit is between 409 and 410KB
+        # We'll set the limit to something in between to be safe
+        if (current_item_size + new_item_size) > 405000:
+            raise ItemSizeTooLarge
+        super(LimitedSizeDict, self).__setitem__(key, value)
+
+    def update(self, *args, **kwargs):
+        if args:
+            if len(args) > 1:
+                raise TypeError("update expected at most 1 arguments, "
+                                "got %d" % len(args))
+            other = dict(args[0])
+            for key in other:
+                self[key] = other[key]
+        for key in kwargs:
+            self[key] = kwargs[key]
+
+    def setdefault(self, key, value=None):
+        if key not in self:
+            self[key] = value
+        return self[key]
+
+
 class Item(BaseModel):
 
     def __init__(self, hash_key, hash_key_type, range_key, range_key_type, attrs):
@@ -134,7 +187,7 @@ class Item(BaseModel):
         self.range_key = range_key
         self.range_key_type = range_key_type
 
-        self.attrs = {}
+        self.attrs = LimitedSizeDict()
         for key, value in attrs.items():
             self.attrs[key] = DynamoType(value)
 

moto/dynamodb2/responses.py

@@ -6,7 +6,7 @@ import re
 
 from moto.core.responses import BaseResponse
 from moto.core.utils import camelcase_to_underscores, amzn_request_id
-from .exceptions import InvalidIndexNameError
+from .exceptions import InvalidIndexNameError, ItemSizeTooLarge
 from .models import dynamodb_backends, dynamo_json_dump
 
 
@@ -225,10 +225,6 @@ class DynamoHandler(BaseResponse):
             er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
             return self.error(er, 'Return values set to invalid value')
 
-        if len(str(item).encode('utf-8')) > 405000:
-            er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
-            return self.error(er, 'Item size has exceeded the maximum allowed size')
-
         if has_empty_keys_or_values(item):
             return get_empty_str_error()
 
@@ -259,6 +255,9 @@ class DynamoHandler(BaseResponse):
                 name, item, expected, condition_expression,
                 expression_attribute_names, expression_attribute_values,
                 overwrite)
+        except ItemSizeTooLarge:
+            er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
+            return self.error(er, ItemSizeTooLarge.message)
         except ValueError:
             er = 'com.amazonaws.dynamodb.v20111205#ConditionalCheckFailedException'
             return self.error(er, 'A condition specified in the operation could not be evaluated.')
@@ -648,6 +647,9 @@ class DynamoHandler(BaseResponse):
                 name, key, update_expression, attribute_updates, expression_attribute_names,
                 expression_attribute_values, expected, condition_expression
             )
+        except ItemSizeTooLarge:
+            er = 'com.amazonaws.dynamodb.v20111205#ValidationException'
+            return self.error(er, ItemSizeTooLarge.message)
         except ValueError:
             er = 'com.amazonaws.dynamodb.v20111205#ConditionalCheckFailedException'
             return self.error(er, 'A condition specified in the operation could not be evaluated.')

tests/test_dynamodb2/test_dynamodb.py

@@ -2282,6 +2282,7 @@ def test_index_with_unknown_attributes_should_fail():
 @mock_dynamodb2
 def test_item_size_is_under_400KB():
     dynamodb = boto3.client('dynamodb', region_name='us-east-1')
+    res = boto3.resource('dynamodb')
 
     dynamodb.create_table(
         TableName='moto-test',
@@ -2289,12 +2290,18 @@ def test_item_size_is_under_400KB():
         AttributeDefinitions=[{'AttributeName': 'id', 'AttributeType': 'S'}],
         ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1}
     )
+    table = res.Table('moto-test')
 
+    large_item = 'x' * 410 * 1000
     with assert_raises(ClientError) as ex:
-        large_item = 'x' * 410 * 1000
+        dynamodb.put_item(TableName='moto-test', Item={'id': {'S': 'foo'}, 'item': {'S': large_item}})
-        dynamodb.put_item(
+    with assert_raises(ClientError) as ex:
-            TableName='moto-test',
+        table.put_item(Item={'id': 'bar', 'item': large_item})
-            Item={'id': {'S': 'foo'}, 'item': {'S': large_item}})
+    with assert_raises(ClientError) as ex:
+        dynamodb.update_item(TableName='moto-test',
+                             Key={'id': {'S': 'foo2'}},
+                             UpdateExpression='set item=:Item',
+                             ExpressionAttributeValues={':Item': {'S': large_item}})
     ex.exception.response['Error']['Code'].should.equal('ValidationException')
     ex.exception.response['Error']['Message'].should.equal('Item size has exceeded the maximum allowed size')