Support for versioned optimistic locking a la DynamoDBMapper (pynamodb#228)

Author: Eric Muller

Diff for: pynamodb/attributes.py

@@ -494,6 +494,13 @@ def deserialize(self, value):
         return json.loads(value)
 
 
+class VersionAttribute(NumberAttribute):
+    """
+    A version attribute
+    """
+    null = True  # should I be doing this?
+
+
 class TTLAttribute(Attribute):
     """
     A time-to-live attribute that signifies when the item expires and can be automatically deleted.

Diff for: pynamodb/attributes.pyi

@@ -119,6 +119,11 @@ class NumberAttribute(Attribute[float]):
     @overload
     def __get__(self, instance: Any, owner: Any) -> float: ...
 
+class VersionAttribute(Attribute[float]):
+    @overload
+    def __get__(self: _A, instance: None, owner: Any) -> _A: ...
+    @overload
+    def __get__(self, instance: Any, owner: Any) -> float: ...
 
 class TTLAttribute(Attribute[datetime]):
     @overload

Diff for: pynamodb/models.py

@@ -10,7 +10,9 @@
 
 from six import add_metaclass
 from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError, InvalidStateError
-from pynamodb.attributes import Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute, TTLAttribute
+from pynamodb.attributes import (
+    Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute, TTLAttribute, VersionAttribute
+)
 from pynamodb.connection.table import TableConnection
 from pynamodb.connection.util import pythonic
 from pynamodb.types import HASH, RANGE
@@ -209,6 +211,14 @@ def __init__(cls, name, bases, attrs):
             if len(ttl_attr_names) > 1:
                 raise ValueError("The model has more than one TTL attribute: {}".format(", ".join(ttl_attr_names)))
 
+            version_attribute_names = [
+                name for name, attr_obj in attrs.items() if isinstance(attr_obj, VersionAttribute)
+            ]
+            if len(version_attribute_names) > 1:
+                raise ValueError(
+                    "The model has more than one Version attribute: {}".format(", ".join(version_attribute_names))
+                )
+
             if META_CLASS_NAME not in attrs:
                 setattr(cls, META_CLASS_NAME, DefaultMeta)
 
@@ -334,6 +344,10 @@ def delete(self, condition=None):
         Deletes this object from dynamodb
         """
         args, kwargs = self._get_save_args(attributes=False, null_check=False)
+        version_condition = self._handle_version_attribute(kwargs)
+        if version_condition is not None:
+            condition &= version_condition
+
         kwargs.update(condition=condition)
         return self._get_connection().delete_item(*args, **kwargs)
 
@@ -348,6 +362,9 @@ def update(self, actions, condition=None):
             raise TypeError("the value of `actions` is expected to be a non-empty list")
 
         args, save_kwargs = self._get_save_args(null_check=False)
+        version_condition = self._handle_version_attribute(save_kwargs, actions=actions)
+        if version_condition is not None:
+            condition &= version_condition
         kwargs = {
             pythonic(RETURN_VALUES):  ALL_NEW,
         }
@@ -371,6 +388,9 @@ def save(self, condition=None):
         Save this object to dynamodb
         """
         args, kwargs = self._get_save_args()
+        version_condition = self._handle_version_attribute(serialized_attributes=kwargs)
+        if version_condition is not None:
+            condition &= version_condition
         kwargs.update(condition=condition)
         return self._get_connection().put_item(*args, **kwargs)
 
@@ -395,6 +415,10 @@ def get_operation_kwargs_from_instance(self,
                                            return_values_on_condition_failure=None):
         is_update = actions is not None
         args, save_kwargs = self._get_save_args(null_check=not is_update)
+        version_condition = self._handle_version_attribute(serialized_attributes=save_kwargs,
+                                                           actions=actions)
+        if version_condition is not None:
+            condition &= version_condition
         kwargs = dict(
             key=key,
             actions=actions,
@@ -872,6 +896,7 @@ def _get_save_args(self, attributes=True, null_check=True):
         """
         kwargs = {}
         serialized = self._serialize(null_check=null_check)
+
         hash_key = serialized.get(HASH)
         range_key = serialized.get(RANGE, None)
         args = (hash_key, )
@@ -881,6 +906,40 @@ def _get_save_args(self, attributes=True, null_check=True):
             kwargs[pythonic(ATTRIBUTES)] = serialized[pythonic(ATTRIBUTES)]
         return args, kwargs
 
+    def _handle_version_attribute(self, serialized_attributes, actions=None):
+        """
+        Handles modifying the request to set or increment the version attribute.
+
+        :param serialized_attributes: A dictionary mapping attribute names to serialized values.
+        :param actions: A non-empty list when performing an update, otherwise None.
+        """
+        version_condition = None
+
+        for name, attr in self.get_attributes().items():
+            value = getattr(self, name)
+            if isinstance(attr, VersionAttribute):
+                # We don't modify the attribute except in the serialized payload so that
+                # the local object is not modified on failure.
+                if not value:
+                    version_condition = attr.does_not_exist()
+                    if actions:
+                        actions.append(attr.set(1))
+                    elif pythonic(ATTRIBUTES) in serialized_attributes:
+                        serialized_attributes[pythonic(ATTRIBUTES)][attr.attr_name] = self._serialize_value(
+                            attr, 1, null_check=True
+                        )
+                else:
+                    version_condition = attr == value
+                    if actions:
+                        actions.append(attr.add(1))
+                    elif pythonic(ATTRIBUTES) in serialized_attributes:
+                        serialized_attributes[pythonic(ATTRIBUTES)][attr.attr_name] = self._serialize_value(
+                            attr, value + 1, null_check=True
+                        )
+
+                break
+        return version_condition
+
     @classmethod
     def _hash_key_attribute(cls):
         """
@@ -1002,7 +1061,6 @@ def _serialize(self, attr_map=False, null_check=True):
             if isinstance(value, MapAttribute):
                 if not value.validate():
                     raise ValueError("Attribute '{}' is not correctly typed".format(attr.attr_name))
-
             serialized = self._serialize_value(attr, value, null_check)
             if NULL in serialized:
                 continue

Diff for: tests/data.py

@@ -1412,3 +1412,38 @@
         "TableStatus": "ACTIVE"
     }
 }
+
+VERSIONED_TABLE_DATA = {
+    "Table": {
+        "AttributeDefinitions": [
+            {
+                "AttributeName": "name",
+                "AttributeType": "S"
+            },
+            {
+                "AttributeName": "email",
+                "AttributeType": "S"
+            },
+            {
+                "AttributeName": "version",
+                "AttributeType": "N"
+            }
+        ],
+        "CreationDateTime": 1.363729002358E9,
+        "ItemCount": 42,
+        "KeySchema": [
+            {
+                "AttributeName": "name",
+                "KeyType": "HASH"
+            },
+        ],
+        "ProvisionedThroughput": {
+            "NumberOfDecreasesToday": 0,
+            "ReadCapacityUnits": 5,
+            "WriteCapacityUnits": 5
+        },
+        "TableName": "VersionedModel",
+        "TableSizeBytes": 0,
+        "TableStatus": "ACTIVE"
+    }
+}

Diff for: tests/test_attributes.py

@@ -18,7 +18,7 @@
     BinarySetAttribute, BinaryAttribute, NumberSetAttribute, NumberAttribute,
     UnicodeAttribute, UnicodeSetAttribute, UTCDateTimeAttribute, BooleanAttribute, MapAttribute,
     ListAttribute, JSONAttribute, TTLAttribute, _get_value_for_deserialize, _fast_parse_utc_datestring,
-)
+    VersionAttribute)
 from pynamodb.constants import (
     DATETIME_FORMAT, DEFAULT_ENCODING, NUMBER, STRING, STRING_SET, NUMBER_SET, BINARY_SET,
     BINARY, BOOLEAN,
@@ -999,3 +999,18 @@ def __eq__(self, other):
         assert deserialized == inp
         assert serialize_mock.call_args_list == [call(1), call(2)]
         assert deserialize_mock.call_args_list == [call('1'), call('2')]
+
+
+class TestVersionAttribute:
+    def test_serialize(self):
+        attr = VersionAttribute()
+        assert attr.attr_type == NUMBER
+        assert attr.serialize(3.141) == '3.141'
+        assert attr.serialize(1) == '1'
+        assert attr.serialize(12345678909876543211234234324234) == '12345678909876543211234234324234'
+
+    def test_deserialize(self):
+        attr = VersionAttribute()
+        assert attr.deserialize('1') == 1
+        assert attr.deserialize('3.141') == 3.141
+        assert attr.deserialize('12345678909876543211234234324234') == 12345678909876543211234234324234