Merge pull request #499 from im-auld/partial-save-support-with-range-key

Partial save support with range key

moto/dynamodb2/models.py

@@ -500,12 +500,20 @@ class DynamoDBBackend(BaseBackend):
     def update_item(self, table_name, key, update_expression, attribute_updates):
         table = self.get_table(table_name)
 
-        if table.hash_key_attr in key:
-            # Sometimes the key is wrapped in a dict with the key name
-            key = key[table.hash_key_attr]
+        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 will be a dict
+            hash_value = DynamoType(key[table.hash_key_attr])
+            range_value = DynamoType(key[table.range_key_attr])
+        elif table.hash_key_attr in key:
+            # Covers tables that have a range key where ``key`` param is a dict
+            hash_value = DynamoType(key[table.hash_key_attr])
+            range_value = None
+        else:
+            # Covers other cases
+            hash_value = DynamoType(key)
+            range_value = None
 
-        hash_value = DynamoType(key)
-        item = table.get_item(hash_value)
+        item = table.get_item(hash_value, range_value)
         if update_expression:
             item.update(update_expression)
         else:

tests/test_dynamodb2/test_dynamodb_table_with_range_key.py

@@ -143,6 +143,36 @@ def test_item_add_and_describe_and_update():
     })
 
 
+@requires_boto_gte("2.9")
+@mock_dynamodb2
+def test_item_partial_save():
+    table = create_table()
+
+    data = {
+        'forum_name': 'LOLCat Forum',
+        'subject': 'The LOLz',
+        'Body': 'http://url_to_lolcat.gif',
+        'SentBy': 'User A',
+    }
+
+    table.put_item(data=data)
+    returned_item = table.get_item(forum_name="LOLCat Forum", subject='The LOLz')
+
+    returned_item['SentBy'] = 'User B'
+    returned_item.partial_save()
+
+    returned_item = table.get_item(
+        forum_name='LOLCat Forum',
+        subject='The LOLz'
+    )
+    dict(returned_item).should.equal({
+        'forum_name': 'LOLCat Forum',
+        'subject': 'The LOLz',
+        'Body': 'http://url_to_lolcat.gif',
+        'SentBy': 'User B',
+    })
+
+
 @requires_boto_gte("2.9")
 @mock_dynamodb2
 def test_item_put_without_table():