Correct behavior of ReturnValues parameter to put_item and update_item

moto/dynamodb2/responses.py

@@ -183,6 +183,7 @@ class DynamoHandler(BaseResponse):
     def put_item(self):
         name = self.body['TableName']
         item = self.body['Item']
+        return_values = self.body.get('ReturnValues', 'NONE')
 
         if has_empty_keys_or_values(item):
             return get_empty_str_error()
@@ -193,6 +194,13 @@ class DynamoHandler(BaseResponse):
         else:
             expected = None
 
+        if return_values == 'ALL_OLD':
+            existing_item = self.dynamodb_backend.get_item(name, item)
+            if existing_item:
+                existing_attributes = existing_item.to_json()['Attributes']
+            else:
+                existing_attributes = {}
+
         # Attempt to parse simple ConditionExpressions into an Expected
         # expression
         if not expected:
@@ -228,6 +236,10 @@ class DynamoHandler(BaseResponse):
                 'TableName': name,
                 'CapacityUnits': 1
             }
+            if return_values == 'ALL_OLD':
+                item_dict['Attributes'] = existing_attributes
+            else:
+                item_dict.pop('Attributes', None)
             return dynamo_json_dump(item_dict)
         else:
             er = 'com.amazonaws.dynamodb.v20111205#ResourceNotFoundException'
@@ -527,9 +539,9 @@ class DynamoHandler(BaseResponse):
         return dynamo_json_dump(item_dict)
 
     def update_item(self):
-
         name = self.body['TableName']
         key = self.body['Key']
+        return_values = self.body.get('ReturnValues', 'NONE')
         update_expression = self.body.get('UpdateExpression')
         attribute_updates = self.body.get('AttributeUpdates')
         expression_attribute_names = self.body.get(
@@ -537,6 +549,10 @@ class DynamoHandler(BaseResponse):
         expression_attribute_values = self.body.get(
             'ExpressionAttributeValues', {})
         existing_item = self.dynamodb_backend.get_item(name, key)
+        if existing_item:
+            existing_attributes = existing_item.to_json()['Attributes']
+        else:
+            existing_attributes = {}
 
         if has_empty_keys_or_values(expression_attribute_values):
             return get_empty_str_error()
@@ -591,8 +607,26 @@ class DynamoHandler(BaseResponse):
             'TableName': name,
             'CapacityUnits': 0.5
         }
-        if not existing_item:
+        unchanged_attributes = {
+            k for k in existing_attributes.keys()
+            if existing_attributes[k] == item_dict['Attributes'].get(k)
+        }
+        changed_attributes = set(existing_attributes.keys()).union(item_dict['Attributes'].keys()).difference(unchanged_attributes)
+
+        if return_values == 'NONE':
             item_dict['Attributes'] = {}
+        elif return_values == 'ALL_OLD':
+            item_dict['Attributes'] = existing_attributes
+        elif return_values == 'UPDATED_OLD':
+            item_dict['Attributes'] = {
+                k: v for k, v in existing_attributes.items()
+                if k in changed_attributes
+            }
+        elif return_values == 'UPDATED_NEW':
+            item_dict['Attributes'] = {
+                k: v for k, v in item_dict['Attributes'].items()
+                if k in changed_attributes
+            }
 
         return dynamo_json_dump(item_dict)
 

tests/test_dynamodb2/test_dynamodb.py

@@ -1246,6 +1246,69 @@ def test_update_if_not_exists():
     assert resp['Items'][0]['created_at'] == 123
 
 
+# https://github.com/spulec/moto/issues/1937
+@mock_dynamodb2
+def test_update_return_attributes():
+    dynamodb = boto3.client('dynamodb', region_name='us-east-1')
+
+    dynamodb.create_table(
+        TableName='moto-test',
+        KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'}],
+        AttributeDefinitions=[{'AttributeName': 'id', 'AttributeType': 'S'}],
+        ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1}
+    )
+
+    def update(col, to, rv):
+        return dynamodb.update_item(
+            TableName='moto-test',
+            Key={'id': {'S': 'foo'}},
+            AttributeUpdates={col: {'Value': {'S': to}, 'Action': 'PUT'}},
+            ReturnValues=rv
+        )
+
+    r = update('col1', 'val1', 'ALL_NEW')
+    assert r['Attributes'] == {'id': {'S': 'foo'}, 'col1': {'S': 'val1'}}
+
+    r = update('col1', 'val2', 'ALL_OLD')
+    assert r['Attributes'] == {'id': {'S': 'foo'}, 'col1': {'S': 'val1'}}
+
+    r = update('col2', 'val3', 'UPDATED_NEW')
+    assert r['Attributes'] == {'col2': {'S': 'val3'}}
+
+    r = update('col2', 'val4', 'UPDATED_OLD')
+    assert r['Attributes'] == {'col2': {'S': 'val3'}}
+
+    r = update('col1', 'val5', 'NONE')
+    assert r['Attributes'] == {}
+
+
+@mock_dynamodb2
+def test_put_return_attributes():
+    dynamodb = boto3.client('dynamodb', region_name='us-east-1')
+
+    dynamodb.create_table(
+        TableName='moto-test',
+        KeySchema=[{'AttributeName': 'id', 'KeyType': 'HASH'}],
+        AttributeDefinitions=[{'AttributeName': 'id', 'AttributeType': 'S'}],
+        ProvisionedThroughput={'ReadCapacityUnits': 1, 'WriteCapacityUnits': 1}
+    )
+
+    r = dynamodb.put_item(
+        TableName='moto-test',
+        Item={'id': {'S': 'foo'}, 'col1': {'S': 'val1'}},
+        ReturnValues='NONE'
+    )
+    assert 'Attributes' not in r
+    
+    r = dynamodb.put_item(
+        TableName='moto-test',
+        Item={'id': {'S': 'foo'}, 'col1': {'S': 'val2'}},
+        ReturnValues='ALL_OLD'
+    )
+    assert r['Attributes'] == {'id': {'S': 'foo'}, 'col1': {'S': 'val1'}}
+    
+    
+
 @mock_dynamodb2
 def test_query_global_secondary_index_when_created_via_update_table_resource():
     dynamodb = boto3.resource('dynamodb', region_name='us-east-1')