Skip to content

Commit 12d8fd2

Browse files
ermulgarrettheel
ermul
authored and
garrettheel
committedAug 5, 2019
Support for versioned optimistic locking a la DynamoDBMapper (#664)
·
stable4.0.0
1 parent 9ccd683 commit 12d8fd2

14 files changed

+759
-19
lines changed
 

‎AUTHORS.rst

+2-1
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,5 @@ PynamoDB is written and maintained by Jharrod LaFon and numerous contributors:
33
* Craig Bruce
44
* Adam Chainz
55
* Andy Wolfe
6-
* Pior Bastida
6+
* Pior Bastida
7+
* Eric Muller

‎docs/optimistic_locking.rst

+172
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
==================
2+
Optimistic Locking
3+
==================
4+
5+
Optimistic Locking is a strategy for ensuring that your database writes are not overwritten by the writes of others.
6+
With optimistic locking, each item has an attribute that acts as a version number. If you retrieve an item from a
7+
table, the application records the version number of that item. You can update the item, but only if the version number
8+
on the server side has not changed. If there is a version mismatch, it means that someone else has modified the item
9+
before you did. The update attempt fails, because you have a stale version of the item. If this happens, you simply
10+
try again by retrieving the item and then trying to update it. Optimistic locking prevents you from accidentally
11+
overwriting changes that were made by others. It also prevents others from accidentally overwriting your changes.
12+
13+
.. warning:: - Optimistic locking will not work properly if you use DynamoDB global tables as they use last-write-wins for concurrent updates.
14+
15+
See also:
16+
`DynamoDBMapper Documentation on Optimistic Locking <https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBMapper.OptimisticLocking.html>`_.
17+
18+
Version Attribute
19+
-----------------
20+
21+
To enable optimistic locking for a table simply add a ``VersionAttribute`` to your model definition.
22+
23+
.. code-block:: python
24+
25+
class OfficeEmployeeMap(MapAttribute):
26+
office_employee_id = UnicodeAttribute()
27+
person = UnicodeAttribute()
28+
29+
def __eq__(self, other):
30+
return isinstance(other, OfficeEmployeeMap) and self.person == other.person
31+
32+
def __repr__(self):
33+
return str(vars(self))
34+
35+
36+
class Office(Model):
37+
class Meta:
38+
read_capacity_units = 1
39+
write_capacity_units = 1
40+
table_name = 'Office'
41+
host = "http://localhost:8000"
42+
office_id = UnicodeAttribute(hash_key=True)
43+
employees = ListAttribute(of=OfficeEmployeeMap)
44+
name = UnicodeAttribute()
45+
version = VersionAttribute()
46+
47+
The attribute is underpinned by an integer which is initialized with 1 when an item is saved for the first time
48+
and is incremented by 1 with each subsequent write operation.
49+
50+
.. code-block:: python
51+
52+
justin = OfficeEmployeeMap(office_employee_id=str(uuid4()), person='justin')
53+
garrett = OfficeEmployeeMap(office_employee_id=str(uuid4()), person='garrett')
54+
office = Office(office_id=str(uuid4()), name="office", employees=[justin, garrett])
55+
office.save()
56+
assert office.version == 1
57+
58+
# Get a second local copy of Office
59+
office_out_of_date = Office.get(office.office_id)
60+
61+
# Add another employee and persist the change.
62+
office.employees.append(OfficeEmployeeMap(office_employee_id=str(uuid4()), person='lita'))
63+
office.save()
64+
# On subsequent save or update operations the version is also incremented locally to match the persisted value so
65+
# there's no need to refresh between operations when reusing the local copy.
66+
assert office.version == 2
67+
assert office_out_of_date.version == 1
68+
69+
The version checking is implemented using DynamoDB conditional write constraints, asserting that no value exists
70+
for the version attribute on the initial save and that the persisted value matches the local value on subsequent writes.
71+
72+
73+
Model.{update, save, delete}
74+
----------------------------
75+
These operations will fail if the local object is out-of-date.
76+
77+
.. code-block:: python
78+
79+
@contextmanager
80+
def assert_condition_check_fails():
81+
try:
82+
yield
83+
except (PutError, UpdateError, DeleteError) as e:
84+
assert isinstance(e.cause, ClientError)
85+
assert e.cause_response_code == "ConditionalCheckFailedException"
86+
except TransactWriteError as e:
87+
assert isinstance(e.cause, ClientError)
88+
assert e.cause_response_code == "TransactionCanceledException"
89+
assert "ConditionalCheckFailed" in e.cause_response_message
90+
else:
91+
raise AssertionError("The version attribute conditional check should have failed.")
92+
93+
94+
with assert_condition_check_fails():
95+
office_out_of_date.update(actions=[Office.name.set('new office name')])
96+
97+
office_out_of_date.employees.remove(garrett)
98+
with assert_condition_check_fails():
99+
office_out_of_date.save()
100+
101+
# After refreshing the local copy our write operations succeed.
102+
office_out_of_date.refresh()
103+
office_out_of_date.employees.remove(garrett)
104+
office_out_of_date.save()
105+
assert office_out_of_date.version == 3
106+
107+
with assert_condition_check_fails():
108+
office.delete()
109+
110+
Transactions
111+
------------
112+
113+
Transactions are supported.
114+
115+
Successful
116+
__________
117+
118+
.. code-block:: python
119+
120+
connection = Connection(host='http://localhost:8000')
121+
122+
office2 = Office(office_id=str(uuid4()), name="second office", employees=[justin])
123+
office2.save()
124+
assert office2.version == 1
125+
office3 = Office(office_id=str(uuid4()), name="third office", employees=[garrett])
126+
office3.save()
127+
assert office3.version == 1
128+
129+
with TransactWrite(connection=connection) as transaction:
130+
transaction.condition_check(Office, office.office_id, condition=(Office.name.exists()))
131+
transaction.delete(office2)
132+
transaction.save(Office(office_id=str(uuid4()), name="new office", employees=[justin, garrett]))
133+
transaction.update(
134+
office3,
135+
actions=[
136+
Office.name.set('birdistheword'),
137+
]
138+
)
139+
140+
try:
141+
office2.refresh()
142+
except DoesNotExist:
143+
pass
144+
else:
145+
raise AssertionError(
146+
'Office with office_id="{}" should have been deleted in the transaction.'
147+
.format(office2.office_id)
148+
)
149+
150+
assert office.version == 2
151+
assert office3.version == 2
152+
153+
Failed
154+
______
155+
156+
.. code-block:: python
157+
158+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
159+
transaction.save(Office(office.office_id, name='newer name', employees=[]))
160+
161+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
162+
transaction.update(
163+
Office(office.office_id, name='newer name', employees=[]),
164+
actions=[Office.name.set('Newer Office Name')]
165+
)
166+
167+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
168+
transaction.delete(Office(office.office_id, name='newer name', employees=[]))
169+
170+
Batch Operations
171+
----------------
172+
*Unsupported* as they do not support conditional writes.

‎examples/optimistic_locking.py

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
from contextlib import contextmanager
2+
from uuid import uuid4
3+
from botocore.client import ClientError
4+
5+
from pynamodb.connection import Connection
6+
from pynamodb.attributes import ListAttribute, MapAttribute, UnicodeAttribute, VersionAttribute
7+
from pynamodb.exceptions import PutError, UpdateError, TransactWriteError, DeleteError, DoesNotExist
8+
from pynamodb.models import Model
9+
from pynamodb.transactions import TransactWrite
10+
11+
12+
class OfficeEmployeeMap(MapAttribute):
13+
office_employee_id = UnicodeAttribute()
14+
person = UnicodeAttribute()
15+
16+
def __eq__(self, other):
17+
return isinstance(other, OfficeEmployeeMap) and self.person == other.person
18+
19+
def __repr__(self):
20+
return str(vars(self))
21+
22+
23+
class Office(Model):
24+
class Meta:
25+
read_capacity_units = 1
26+
write_capacity_units = 1
27+
table_name = 'Office'
28+
host = "http://localhost:8000"
29+
office_id = UnicodeAttribute(hash_key=True)
30+
employees = ListAttribute(of=OfficeEmployeeMap)
31+
name = UnicodeAttribute()
32+
version = VersionAttribute()
33+
34+
35+
if not Office.exists():
36+
Office.create_table(wait=True)
37+
38+
39+
@contextmanager
40+
def assert_condition_check_fails():
41+
try:
42+
yield
43+
except (PutError, UpdateError, DeleteError) as e:
44+
assert isinstance(e.cause, ClientError)
45+
assert e.cause_response_code == "ConditionalCheckFailedException"
46+
except TransactWriteError as e:
47+
assert isinstance(e.cause, ClientError)
48+
assert e.cause_response_code == "TransactionCanceledException"
49+
assert "ConditionalCheckFailed" in e.cause_response_message
50+
else:
51+
raise AssertionError("The version attribute conditional check should have failed.")
52+
53+
54+
justin = OfficeEmployeeMap(office_employee_id=str(uuid4()), person='justin')
55+
garrett = OfficeEmployeeMap(office_employee_id=str(uuid4()), person='garrett')
56+
office = Office(office_id=str(uuid4()), name="office 3", employees=[justin, garrett])
57+
office.save()
58+
assert office.version == 1
59+
60+
# Get a second local copy of Office
61+
office_out_of_date = Office.get(office.office_id)
62+
# Add another employee and save the changes.
63+
office.employees.append(OfficeEmployeeMap(office_employee_id=str(uuid4()), person='lita'))
64+
office.save()
65+
# After a successful save or update operation the version is set or incremented locally so there's no need to refresh
66+
# between operations using the same local copy.
67+
assert office.version == 2
68+
assert office_out_of_date.version == 1
69+
70+
# Condition check fails for update.
71+
with assert_condition_check_fails():
72+
office_out_of_date.update(actions=[Office.name.set('new office name')])
73+
74+
# Condition check fails for save.
75+
office_out_of_date.employees.remove(garrett)
76+
with assert_condition_check_fails():
77+
office_out_of_date.save()
78+
79+
# After refreshing the local copy the operation will succeed.
80+
office_out_of_date.refresh()
81+
office_out_of_date.employees.remove(garrett)
82+
office_out_of_date.save()
83+
assert office_out_of_date.version == 3
84+
85+
# Condition check fails for delete.
86+
with assert_condition_check_fails():
87+
office.delete()
88+
89+
# Example failed transactions.
90+
connection = Connection(host='http://localhost:8000')
91+
92+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
93+
transaction.save(Office(office.office_id, name='newer name', employees=[]))
94+
95+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
96+
transaction.update(
97+
Office(office.office_id, name='newer name', employees=[]),
98+
actions=[
99+
Office.name.set('Newer Office Name'),
100+
]
101+
)
102+
103+
with assert_condition_check_fails(), TransactWrite(connection=connection) as transaction:
104+
transaction.delete(Office(office.office_id, name='newer name', employees=[]))
105+
106+
# Example successful transaction.
107+
office2 = Office(office_id=str(uuid4()), name="second office", employees=[justin])
108+
office2.save()
109+
assert office2.version == 1
110+
office3 = Office(office_id=str(uuid4()), name="third office", employees=[garrett])
111+
office3.save()
112+
assert office3.version == 1
113+
114+
with TransactWrite(connection=connection) as transaction:
115+
transaction.condition_check(Office, office.office_id, condition=(Office.name.exists()))
116+
transaction.delete(office2)
117+
transaction.save(Office(office_id=str(uuid4()), name="new office", employees=[justin, garrett]))
118+
transaction.update(
119+
office3,
120+
actions=[
121+
Office.name.set('birdistheword'),
122+
]
123+
)
124+
125+
try:
126+
office2.refresh()
127+
except DoesNotExist:
128+
pass
129+
else:
130+
raise AssertionError(
131+
"This item should have been deleted, but no DoesNotExist "
132+
"exception was raised when attempting to refresh a local copy."
133+
)
134+
135+
assert office.version == 2
136+
# The version attribute of items which are saved or updated in a transaction are updated automatically to match the
137+
# persisted value.
138+
assert office3.version == 2
139+
office.refresh()
140+
assert office.version == 3

‎pynamodb/attributes.py

+32
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,38 @@ def deserialize(self, value):
494494
return json.loads(value)
495495

496496

497+
class VersionAttribute(NumberAttribute):
498+
"""
499+
A version attribute
500+
"""
501+
null = True
502+
503+
def __set__(self, instance, value):
504+
"""
505+
Cast assigned value to int.
506+
"""
507+
super(VersionAttribute, self).__set__(instance, int(value))
508+
509+
def __get__(self, instance, owner):
510+
"""
511+
Cast retrieved value to int.
512+
"""
513+
val = super(VersionAttribute, self).__get__(instance, owner)
514+
return int(val) if isinstance(val, float) else val
515+
516+
def serialize(self, value):
517+
"""
518+
Cast value to int then encode as JSON
519+
"""
520+
return super(VersionAttribute, self).serialize(int(value))
521+
522+
def deserialize(self, value):
523+
"""
524+
Decode numbers from JSON and cast to int.
525+
"""
526+
return int(super(VersionAttribute, self).deserialize(value))
527+
528+
497529
class TTLAttribute(Attribute):
498530
"""
499531
A time-to-live attribute that signifies when the item expires and can be automatically deleted.

‎pynamodb/attributes.pyi

+5
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,11 @@ class NumberAttribute(Attribute[float]):
119119
@overload
120120
def __get__(self, instance: Any, owner: Any) -> float: ...
121121

122+
class VersionAttribute(NumberAttribute):
123+
@overload
124+
def __get__(self: _A, instance: None, owner: Any) -> _A: ...
125+
@overload
126+
def __get__(self, instance: Any, owner: Any) -> int: ...
122127

123128
class TTLAttribute(Attribute[datetime]):
124129
@overload

‎pynamodb/connection/base.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -849,7 +849,7 @@ def get_operation_kwargs(self,
849849

850850
operation_kwargs[TABLE_NAME] = table_name
851851
operation_kwargs.update(self.get_identifier_map(table_name, hash_key, range_key, key=key))
852-
if attributes:
852+
if attributes and operation_kwargs.get(ITEM) is not None:
853853
attrs = self.get_item_attribute_map(table_name, attributes)
854854
operation_kwargs[ITEM].update(attrs[ITEM])
855855
if attributes_to_get is not None:

‎pynamodb/models.py

+72-2
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,12 @@
99
from inspect import getmembers
1010

1111
from six import add_metaclass
12+
13+
from pynamodb.expressions.condition import NotExists, Comparison
1214
from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError, InvalidStateError
13-
from pynamodb.attributes import Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute, TTLAttribute
15+
from pynamodb.attributes import (
16+
Attribute, AttributeContainer, AttributeContainerMeta, MapAttribute, TTLAttribute, VersionAttribute
17+
)
1418
from pynamodb.connection.table import TableConnection
1519
from pynamodb.connection.util import pythonic
1620
from pynamodb.types import HASH, RANGE
@@ -172,6 +176,13 @@ def __init__(cls, name, bases, attrs):
172176
cls._hash_keyname = attr_name
173177
if attribute.is_range_key:
174178
cls._range_keyname = attr_name
179+
if isinstance(attribute, VersionAttribute):
180+
if cls._version_attribute_name:
181+
raise ValueError(
182+
"The model has more than one Version attribute: {}, {}"
183+
.format(cls._version_attribute_name, attr_name)
184+
)
185+
cls._version_attribute_name = attr_name
175186
if isinstance(attrs, dict):
176187
for attr_name, attr_obj in attrs.items():
177188
if attr_name == META_CLASS_NAME:
@@ -238,6 +249,7 @@ class Model(AttributeContainer):
238249
_connection = None
239250
_index_classes = None
240251
DoesNotExist = DoesNotExist
252+
_version_attribute_name = None
241253

242254
def __init__(self, hash_key=None, range_key=None, _user_instantiated=True, **attributes):
243255
"""
@@ -334,6 +346,10 @@ def delete(self, condition=None):
334346
Deletes this object from dynamodb
335347
"""
336348
args, kwargs = self._get_save_args(attributes=False, null_check=False)
349+
version_condition = self._handle_version_attribute(kwargs)
350+
if version_condition is not None:
351+
condition &= version_condition
352+
337353
kwargs.update(condition=condition)
338354
return self._get_connection().delete_item(*args, **kwargs)
339355

@@ -348,6 +364,9 @@ def update(self, actions, condition=None):
348364
raise TypeError("the value of `actions` is expected to be a non-empty list")
349365

350366
args, save_kwargs = self._get_save_args(null_check=False)
367+
version_condition = self._handle_version_attribute(save_kwargs, actions=actions)
368+
if version_condition is not None:
369+
condition &= version_condition
351370
kwargs = {
352371
pythonic(RETURN_VALUES): ALL_NEW,
353372
}
@@ -371,8 +390,13 @@ def save(self, condition=None):
371390
Save this object to dynamodb
372391
"""
373392
args, kwargs = self._get_save_args()
393+
version_condition = self._handle_version_attribute(serialized_attributes=kwargs)
394+
if version_condition is not None:
395+
condition &= version_condition
374396
kwargs.update(condition=condition)
375-
return self._get_connection().put_item(*args, **kwargs)
397+
data = self._get_connection().put_item(*args, **kwargs)
398+
self.update_local_version_attribute()
399+
return data
376400

377401
def refresh(self, consistent_read=False):
378402
"""
@@ -394,7 +418,16 @@ def get_operation_kwargs_from_instance(self,
394418
condition=None,
395419
return_values_on_condition_failure=None):
396420
is_update = actions is not None
421+
is_delete = actions is None and key is KEY
397422
args, save_kwargs = self._get_save_args(null_check=not is_update)
423+
424+
version_condition = self._handle_version_attribute(
425+
serialized_attributes={} if is_delete else save_kwargs,
426+
actions=actions
427+
)
428+
if version_condition is not None:
429+
condition &= version_condition
430+
398431
kwargs = dict(
399432
key=key,
400433
actions=actions,
@@ -881,6 +914,43 @@ def _get_save_args(self, attributes=True, null_check=True):
881914
kwargs[pythonic(ATTRIBUTES)] = serialized[pythonic(ATTRIBUTES)]
882915
return args, kwargs
883916

917+
def _handle_version_attribute(self, serialized_attributes, actions=None):
918+
"""
919+
Handles modifying the request to set or increment the version attribute.
920+
921+
:param serialized_attributes: A dictionary mapping attribute names to serialized values.
922+
:param actions: A non-empty list when performing an update, otherwise None.
923+
"""
924+
if self._version_attribute_name is None:
925+
return
926+
927+
version_attribute = self.get_attributes()[self._version_attribute_name]
928+
version_attribute_value = getattr(self, self._version_attribute_name)
929+
930+
if version_attribute_value:
931+
version_condition = version_attribute == version_attribute_value
932+
if actions:
933+
actions.append(version_attribute.add(1))
934+
elif pythonic(ATTRIBUTES) in serialized_attributes:
935+
serialized_attributes[pythonic(ATTRIBUTES)][version_attribute.attr_name] = self._serialize_value(
936+
version_attribute, version_attribute_value + 1, null_check=True
937+
)
938+
else:
939+
version_condition = version_attribute.does_not_exist()
940+
if actions:
941+
actions.append(version_attribute.set(1))
942+
elif pythonic(ATTRIBUTES) in serialized_attributes:
943+
serialized_attributes[pythonic(ATTRIBUTES)][version_attribute.attr_name] = self._serialize_value(
944+
version_attribute, 1, null_check=True
945+
)
946+
947+
return version_condition
948+
949+
def update_local_version_attribute(self):
950+
if self._version_attribute_name:
951+
value = getattr(self, self._version_attribute_name, None) or 0
952+
setattr(self, self._version_attribute_name, value + 1)
953+
884954
@classmethod
885955
def _hash_key_attribute(cls):
886956
"""

‎pynamodb/transactions.py

+7-1
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ def __init__(self, client_request_token=None, return_item_collection_metrics=Non
7272
self._delete_items = []
7373
self._put_items = []
7474
self._update_items = []
75+
self._models_for_version_attribute_update = []
7576

7677
def condition_check(self, model_cls, hash_key, range_key=None, condition=None):
7778
if condition is None:
@@ -94,6 +95,7 @@ def save(self, model, condition=None, return_values=None):
9495
return_values_on_condition_failure=return_values
9596
)
9697
self._put_items.append(operation_kwargs)
98+
self._models_for_version_attribute_update.append(model)
9799

98100
def update(self, model, actions, condition=None, return_values=None):
99101
operation_kwargs = model.get_operation_kwargs_from_instance(
@@ -102,9 +104,10 @@ def update(self, model, actions, condition=None, return_values=None):
102104
return_values_on_condition_failure=return_values
103105
)
104106
self._update_items.append(operation_kwargs)
107+
self._models_for_version_attribute_update.append(model)
105108

106109
def _commit(self):
107-
return self._connection.transact_write_items(
110+
response = self._connection.transact_write_items(
108111
condition_check_items=self._condition_check_items,
109112
delete_items=self._delete_items,
110113
put_items=self._put_items,
@@ -113,3 +116,6 @@ def _commit(self):
113116
return_consumed_capacity=self._return_consumed_capacity,
114117
return_item_collection_metrics=self._return_item_collection_metrics,
115118
)
119+
for model in self._models_for_version_attribute_update:
120+
model.update_local_version_attribute()
121+
return response

‎tests/data.py

+35
Original file line numberDiff line numberDiff line change
@@ -1412,3 +1412,38 @@
14121412
"TableStatus": "ACTIVE"
14131413
}
14141414
}
1415+
1416+
VERSIONED_TABLE_DATA = {
1417+
"Table": {
1418+
"AttributeDefinitions": [
1419+
{
1420+
"AttributeName": "name",
1421+
"AttributeType": "S"
1422+
},
1423+
{
1424+
"AttributeName": "email",
1425+
"AttributeType": "S"
1426+
},
1427+
{
1428+
"AttributeName": "version",
1429+
"AttributeType": "N"
1430+
}
1431+
],
1432+
"CreationDateTime": 1.363729002358E9,
1433+
"ItemCount": 42,
1434+
"KeySchema": [
1435+
{
1436+
"AttributeName": "name",
1437+
"KeyType": "HASH"
1438+
},
1439+
],
1440+
"ProvisionedThroughput": {
1441+
"NumberOfDecreasesToday": 0,
1442+
"ReadCapacityUnits": 5,
1443+
"WriteCapacityUnits": 5
1444+
},
1445+
"TableName": "VersionedModel",
1446+
"TableSizeBytes": 0,
1447+
"TableStatus": "ACTIVE"
1448+
}
1449+
}

‎tests/integration/model_integration_test.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55
from pynamodb.models import Model
66
from pynamodb.indexes import GlobalSecondaryIndex, AllProjection, LocalSecondaryIndex
77
from pynamodb.attributes import (
8-
UnicodeAttribute, BinaryAttribute, UTCDateTimeAttribute, NumberSetAttribute, NumberAttribute
9-
)
8+
UnicodeAttribute, BinaryAttribute, UTCDateTimeAttribute, NumberSetAttribute, NumberAttribute,
9+
VersionAttribute)
1010

1111
import pytest
1212

@@ -51,6 +51,7 @@ class Meta:
5151
epoch = UTCDateTimeAttribute(default=datetime.now)
5252
content = BinaryAttribute(null=True)
5353
scores = NumberSetAttribute()
54+
version = VersionAttribute()
5455

5556
if not TestModel.exists():
5657
print("Creating table")
@@ -100,3 +101,4 @@ class Meta:
100101
print("Item queried from index: {}".format(item.view))
101102

102103
print(query_obj.update([TestModel.view.add(1)], condition=TestModel.forum.exists()))
104+
TestModel.delete_table()

‎tests/integration/test_transaction_integration.py

+92-1
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,10 @@
66
from pynamodb.connection import Connection
77
from pynamodb.exceptions import DoesNotExist, TransactWriteError, TransactGetError, InvalidStateError
88

9-
from pynamodb.attributes import NumberAttribute, UnicodeAttribute, UTCDateTimeAttribute, BooleanAttribute
9+
10+
from pynamodb.attributes import (
11+
NumberAttribute, UnicodeAttribute, UTCDateTimeAttribute, BooleanAttribute, VersionAttribute
12+
)
1013
from pynamodb.transactions import TransactGet, TransactWrite
1114

1215
from pynamodb.models import Model
@@ -59,11 +62,22 @@ class Meta:
5962
entry_index = NumberAttribute(hash_key=True)
6063

6164

65+
class Foo(Model):
66+
class Meta:
67+
region = 'us-east-1'
68+
table_name = 'foo'
69+
70+
bar = NumberAttribute(hash_key=True)
71+
star = UnicodeAttribute(null=True)
72+
version = VersionAttribute()
73+
74+
6275
TEST_MODELS = [
6376
BankStatement,
6477
DifferentRegion,
6578
LineItem,
6679
User,
80+
Foo
6781
]
6882

6983

@@ -271,3 +285,80 @@ def test_transact_write__one_of_each(connection):
271285
statement.refresh()
272286
assert not statement.active
273287
assert statement.balance == 0
288+
289+
290+
@pytest.mark.ddblocal
291+
def test_transaction_write_with_version_attribute(connection):
292+
foo1 = Foo(1)
293+
foo1.save()
294+
foo2 = Foo(2, star='bar')
295+
foo2.save()
296+
foo3 = Foo(3)
297+
foo3.save()
298+
299+
with TransactWrite(connection=connection) as transaction:
300+
transaction.condition_check(Foo, 1, condition=(Foo.bar.exists()))
301+
transaction.delete(foo2)
302+
transaction.save(Foo(4))
303+
transaction.update(
304+
foo3,
305+
actions=[
306+
Foo.star.set('birdistheword'),
307+
]
308+
)
309+
310+
assert Foo.get(1).version == 1
311+
with pytest.raises(DoesNotExist):
312+
Foo.get(2)
313+
# Local object's version attribute is updated automatically.
314+
assert foo3.version == 2
315+
assert Foo.get(4).version == 1
316+
317+
318+
@pytest.mark.ddblocal
319+
def test_transaction_get_with_version_attribute(connection):
320+
Foo(11).save()
321+
Foo(12, star='bar').save()
322+
323+
with TransactGet(connection=connection) as transaction:
324+
foo1_future = transaction.get(Foo, 11)
325+
foo2_future = transaction.get(Foo, 12)
326+
327+
foo1 = foo1_future.get()
328+
assert foo1.version == 1
329+
foo2 = foo2_future.get()
330+
assert foo2.version == 1
331+
assert foo2.star == 'bar'
332+
333+
334+
@pytest.mark.ddblocal
335+
def test_transaction_write_with_version_attribute_condition_failure(connection):
336+
foo = Foo(21)
337+
foo.save()
338+
339+
foo2 = Foo(21)
340+
341+
with pytest.raises(TransactWriteError) as exc_info:
342+
with TransactWrite(connection=connection) as transaction:
343+
transaction.save(Foo(21))
344+
assert get_error_code(exc_info.value) == TRANSACTION_CANCELLED
345+
assert 'ConditionalCheckFailed' in get_error_message(exc_info.value)
346+
347+
with pytest.raises(TransactWriteError) as exc_info:
348+
with TransactWrite(connection=connection) as transaction:
349+
transaction.update(
350+
foo2,
351+
actions=[
352+
Foo.star.set('birdistheword'),
353+
]
354+
)
355+
assert get_error_code(exc_info.value) == TRANSACTION_CANCELLED
356+
assert 'ConditionalCheckFailed' in get_error_message(exc_info.value)
357+
# Version attribute is not updated on failure.
358+
assert foo2.version is None
359+
360+
with pytest.raises(TransactWriteError) as exc_info:
361+
with TransactWrite(connection=connection) as transaction:
362+
transaction.delete(foo2)
363+
assert get_error_code(exc_info.value) == TRANSACTION_CANCELLED
364+
assert 'ConditionalCheckFailed' in get_error_message(exc_info.value)

‎tests/test_attributes.py

+16-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
BinarySetAttribute, BinaryAttribute, NumberSetAttribute, NumberAttribute,
1919
UnicodeAttribute, UnicodeSetAttribute, UTCDateTimeAttribute, BooleanAttribute, MapAttribute,
2020
ListAttribute, JSONAttribute, TTLAttribute, _get_value_for_deserialize, _fast_parse_utc_datestring,
21-
)
21+
VersionAttribute)
2222
from pynamodb.constants import (
2323
DATETIME_FORMAT, DEFAULT_ENCODING, NUMBER, STRING, STRING_SET, NUMBER_SET, BINARY_SET,
2424
BINARY, BOOLEAN,
@@ -1004,3 +1004,18 @@ def __eq__(self, other):
10041004
assert deserialized == inp
10051005
assert serialize_mock.call_args_list == [call(1), call(2)]
10061006
assert deserialize_mock.call_args_list == [call('1'), call('2')]
1007+
1008+
1009+
class TestVersionAttribute:
1010+
def test_serialize(self):
1011+
attr = VersionAttribute()
1012+
assert attr.attr_type == NUMBER
1013+
assert attr.serialize(3.141) == '3'
1014+
assert attr.serialize(1) == '1'
1015+
assert attr.serialize(12345678909876543211234234324234) == '12345678909876543211234234324234'
1016+
1017+
def test_deserialize(self):
1018+
attr = VersionAttribute()
1019+
assert attr.deserialize('1') == 1
1020+
assert attr.deserialize('3.141') == 3
1021+
assert attr.deserialize('12345678909876543211234234324234') == 12345678909876543211234234324234

‎tests/test_model.py

+170-4
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
from pynamodb.attributes import (
3131
UnicodeAttribute, NumberAttribute, BinaryAttribute, UTCDateTimeAttribute,
3232
UnicodeSetAttribute, NumberSetAttribute, BinarySetAttribute, MapAttribute,
33-
BooleanAttribute, ListAttribute, TTLAttribute)
33+
BooleanAttribute, ListAttribute, TTLAttribute, VersionAttribute)
3434
from .data import (
3535
MODEL_TABLE_DATA, GET_MODEL_ITEM_DATA, SIMPLE_MODEL_TABLE_DATA,
3636
DESCRIBE_TABLE_DATA_PAY_PER_REQUEST,
@@ -41,12 +41,13 @@
4141
GET_OFFICE_EMPLOYEE_ITEM_DATA, GET_OFFICE_EMPLOYEE_ITEM_DATA_WITH_NULL,
4242
GROCERY_LIST_MODEL_TABLE_DATA, GET_GROCERY_LIST_ITEM_DATA,
4343
GET_OFFICE_ITEM_DATA, OFFICE_MODEL_TABLE_DATA, COMPLEX_MODEL_TABLE_DATA, COMPLEX_MODEL_ITEM_DATA,
44-
CAR_MODEL_TABLE_DATA, FULL_CAR_MODEL_ITEM_DATA, CAR_MODEL_WITH_NULL_ITEM_DATA, INVALID_CAR_MODEL_WITH_NULL_ITEM_DATA,
44+
CAR_MODEL_TABLE_DATA, FULL_CAR_MODEL_ITEM_DATA, CAR_MODEL_WITH_NULL_ITEM_DATA,
45+
INVALID_CAR_MODEL_WITH_NULL_ITEM_DATA,
4546
BOOLEAN_MODEL_TABLE_DATA, BOOLEAN_MODEL_FALSE_ITEM_DATA, BOOLEAN_MODEL_TRUE_ITEM_DATA,
4647
TREE_MODEL_TABLE_DATA, TREE_MODEL_ITEM_DATA,
4748
EXPLICIT_RAW_MAP_MODEL_TABLE_DATA, EXPLICIT_RAW_MAP_MODEL_ITEM_DATA,
48-
EXPLICIT_RAW_MAP_MODEL_AS_SUB_MAP_IN_TYPED_MAP_ITEM_DATA, EXPLICIT_RAW_MAP_MODEL_AS_SUB_MAP_IN_TYPED_MAP_TABLE_DATA
49-
)
49+
EXPLICIT_RAW_MAP_MODEL_AS_SUB_MAP_IN_TYPED_MAP_ITEM_DATA, EXPLICIT_RAW_MAP_MODEL_AS_SUB_MAP_IN_TYPED_MAP_TABLE_DATA,
50+
VERSIONED_TABLE_DATA)
5051

5152
if six.PY3:
5253
from unittest.mock import patch, MagicMock
@@ -440,6 +441,15 @@ class Meta:
440441
my_ttl = TTLAttribute(default_for_new=timedelta(minutes=1))
441442

442443

444+
class VersionedModel(Model):
445+
class Meta:
446+
table_name = 'VersionedModel'
447+
448+
name = UnicodeAttribute(hash_key=True)
449+
email = UnicodeAttribute()
450+
version = VersionAttribute()
451+
452+
443453
class ModelTestCase(TestCase):
444454
"""
445455
Tests for the models API
@@ -2964,6 +2974,153 @@ def fake_dynamodb(*args, **kwargs):
29642974
self.assert_dict_lists_equal(actual['AttributeDefinitions'],
29652975
DOG_TABLE_DATA['Table']['AttributeDefinitions'])
29662976

2977+
def test_model_version_attribute_save(self):
2978+
self.init_table_meta(VersionedModel, VERSIONED_TABLE_DATA)
2979+
item = VersionedModel('test_user_name', email='test_user@email.com')
2980+
2981+
with patch(PATCH_METHOD) as req:
2982+
req.return_value = {}
2983+
item.save()
2984+
args = req.call_args[0][1]
2985+
params = {
2986+
'Item': {
2987+
'name': {
2988+
'S': 'test_user_name'
2989+
},
2990+
'email': {
2991+
'S': 'test_user@email.com'
2992+
},
2993+
'version': {
2994+
'N': '1'
2995+
},
2996+
},
2997+
'ReturnConsumedCapacity': 'TOTAL',
2998+
'TableName': 'VersionedModel',
2999+
'ConditionExpression': 'attribute_not_exists (#0)',
3000+
'ExpressionAttributeNames': {'#0': 'version'},
3001+
}
3002+
3003+
deep_eq(args, params, _assert=True)
3004+
item.version = 1
3005+
item.name = "test_new_username"
3006+
item.save()
3007+
args = req.call_args[0][1]
3008+
3009+
params = {
3010+
'Item': {
3011+
'name': {
3012+
'S': 'test_new_username'
3013+
},
3014+
'email': {
3015+
'S': 'test_user@email.com'
3016+
},
3017+
'version': {
3018+
'N': '2'
3019+
},
3020+
},
3021+
'ReturnConsumedCapacity': 'TOTAL',
3022+
'TableName': 'VersionedModel',
3023+
'ConditionExpression': '#0 = :0',
3024+
'ExpressionAttributeNames': {'#0': 'version'},
3025+
'ExpressionAttributeValues': {':0': {'N': '1'}}
3026+
}
3027+
3028+
deep_eq(args, params, _assert=True)
3029+
3030+
def test_version_attribute_increments_on_update(self):
3031+
self.init_table_meta(VersionedModel, VERSIONED_TABLE_DATA)
3032+
item = VersionedModel('test_user_name', email='test_user@email.com')
3033+
3034+
with patch(PATCH_METHOD) as req:
3035+
req.return_value = {
3036+
ATTRIBUTES: {
3037+
'name': {
3038+
'S': 'test_user_name'
3039+
},
3040+
'email': {
3041+
'S': 'new@email.com'
3042+
},
3043+
'version': {
3044+
'N': '1'
3045+
},
3046+
}
3047+
}
3048+
item.update(actions=[VersionedModel.email.set('new@email.com')])
3049+
args = req.call_args[0][1]
3050+
params = {
3051+
'ConditionExpression': 'attribute_not_exists (#0)',
3052+
'ExpressionAttributeNames': {
3053+
'#0': 'version',
3054+
'#1': 'email'
3055+
},
3056+
'ExpressionAttributeValues': {
3057+
':0': {
3058+
'S': 'new@email.com'
3059+
},
3060+
':1': {
3061+
'N': '1'
3062+
}
3063+
},
3064+
'Key': {
3065+
'name': {
3066+
'S': 'test_user_name'
3067+
}
3068+
},
3069+
'ReturnConsumedCapacity': 'TOTAL',
3070+
'ReturnValues': 'ALL_NEW',
3071+
'TableName': 'VersionedModel',
3072+
'UpdateExpression': 'SET #1 = :0, #0 = :1'
3073+
}
3074+
3075+
deep_eq(args, params, _assert=True)
3076+
assert item.version == 1
3077+
3078+
req.return_value = {
3079+
ATTRIBUTES: {
3080+
'name': {
3081+
'S': 'test_user_name'
3082+
},
3083+
'email': {
3084+
'S': 'newer@email.com'
3085+
},
3086+
'version': {
3087+
'N': '2'
3088+
},
3089+
}
3090+
}
3091+
3092+
item.update(actions=[VersionedModel.email.set('newer@email.com')])
3093+
args = req.call_args[0][1]
3094+
params = {
3095+
'ConditionExpression': '#0 = :0',
3096+
'ExpressionAttributeNames': {
3097+
'#0': 'version',
3098+
'#1': 'email'
3099+
},
3100+
'ExpressionAttributeValues': {
3101+
':0': {
3102+
'N': '1'
3103+
},
3104+
':1': {
3105+
'S': 'newer@email.com'
3106+
},
3107+
':2': {
3108+
'N': '1'
3109+
}
3110+
},
3111+
'Key': {
3112+
'name': {
3113+
'S': 'test_user_name'
3114+
}
3115+
},
3116+
'ReturnConsumedCapacity': 'TOTAL',
3117+
'ReturnValues': 'ALL_NEW',
3118+
'TableName': 'VersionedModel',
3119+
'UpdateExpression': 'SET #1 = :1 ADD #0 :2'
3120+
}
3121+
3122+
deep_eq(args, params, _assert=True)
3123+
29673124

29683125
class ModelInitTestCase(TestCase):
29693126

@@ -3075,3 +3232,12 @@ def test_deserialized_with_ttl(self):
30753232
req.return_value = SIMPLE_MODEL_TABLE_DATA
30763233
m = TTLModel.from_raw_data({'user_name': {'S': 'mock'}, 'my_ttl': {'N': '1546300800'}})
30773234
assert m.my_ttl == datetime(2019, 1, 1, tzinfo=tzutc())
3235+
3236+
def test_multiple_version_attributes(self):
3237+
with self.assertRaises(ValueError):
3238+
class BadVersionedModel(Model):
3239+
class Meta:
3240+
table_name = 'BadVersionedModel'
3241+
3242+
version = VersionAttribute()
3243+
another_version = VersionAttribute()

‎tests/test_transaction.py

+11-6
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import pytest
22
import six
3-
from pynamodb.attributes import NumberAttribute, UnicodeAttribute
3+
from pynamodb.attributes import NumberAttribute, UnicodeAttribute, VersionAttribute
44

55
from pynamodb.connection import Connection
66
from pynamodb.transactions import Transaction, TransactGet, TransactWrite
@@ -20,6 +20,7 @@ class Meta:
2020
mock_hash = NumberAttribute(hash_key=True)
2121
mock_range = NumberAttribute(range_key=True)
2222
mock_toot = UnicodeAttribute(null=True)
23+
mock_version = VersionAttribute()
2324

2425

2526
MOCK_TABLE_DESCRIPTOR = {
@@ -99,22 +100,26 @@ def test_commit(self, mocker):
99100
'TableName': 'mock'}
100101
]
101102
expected_deletes = [{
103+
'ConditionExpression': 'attribute_not_exists (#0)',
104+
'ExpressionAttributeNames': {'#0': 'mock_version'},
102105
'Key': {'MockHash': {'N': '2'}, 'MockRange': {'N': '4'}},
103106
'TableName': 'mock'
104107
}]
105108
expected_puts = [{
106-
'Item': {'MockHash': {'N': '3'}, 'MockRange': {'N': '5'}},
109+
'ConditionExpression': 'attribute_not_exists (#0)',
110+
'ExpressionAttributeNames': {'#0': 'mock_version'},
111+
'Item': {'MockHash': {'N': '3'}, 'MockRange': {'N': '5'}, 'mock_version': {'N': '1'}},
107112
'TableName': 'mock'
108113
}]
109114
expected_updates = [{
115+
'ConditionExpression': 'attribute_not_exists (#0)',
110116
'TableName': 'mock',
111117
'Key': {'MockHash': {'N': '4'}, 'MockRange': {'N': '6'}},
112118
'ReturnValuesOnConditionCheckFailure': 'ALL_OLD',
113-
'UpdateExpression': 'SET #0 = :0',
114-
'ExpressionAttributeNames': {'#0': 'mock_toot'},
115-
'ExpressionAttributeValues': {':0': {'S': 'hello'}}
119+
'UpdateExpression': 'SET #1 = :0, #0 = :1',
120+
'ExpressionAttributeNames': {'#0': 'mock_version', '#1': 'mock_toot'},
121+
'ExpressionAttributeValues': {':0': {'S': 'hello'}, ':1': {'N': '1'}}
116122
}]
117-
118123
mock_connection_transact_write.assert_called_once_with(
119124
condition_check_items=expected_condition_checks,
120125
delete_items=expected_deletes,

0 commit comments

Comments
 (0)
Please sign in to comment.