diff --git a/pynamodb/__init__.py b/pynamodb/__init__.py index 8cc90a6ef..747d3084a 100644 --- a/pynamodb/__init__.py +++ b/pynamodb/__init__.py @@ -7,4 +7,4 @@ """ __author__ = 'Jharrod LaFon' __license__ = 'MIT' -__version__ = '2.1.6' +__version__ = '2.2.0' diff --git a/pynamodb/attributes.py b/pynamodb/attributes.py index 333ae0baa..a7945519f 100644 --- a/pynamodb/attributes.py +++ b/pynamodb/attributes.py @@ -530,6 +530,18 @@ def _get_deserialize_class(cls, key, value): return cls._get_attributes().get(key) return _get_class_for_deserialize(value) + @classmethod + def _has_unicode_set_attribute(cls): + if cls.is_raw(): + return False + + for attr_value in cls._get_attributes().values(): + if isinstance(attr_value, UnicodeSetAttribute): + return True + if isinstance(attr_value, MapAttribute): + return attr_value._has_unicode_set_attribute() + + return False def _get_value_for_deserialize(value): key = next(iter(value.keys())) diff --git a/pynamodb/models.py b/pynamodb/models.py index 81bdfa83d..630462b64 100644 --- a/pynamodb/models.py +++ b/pynamodb/models.py @@ -11,7 +11,7 @@ from six import add_metaclass from pynamodb.exceptions import DoesNotExist, TableDoesNotExist, TableError from pynamodb.throttle import NoThrottle -from pynamodb.attributes import Attribute, AttributeContainer, MapAttribute, ListAttribute +from pynamodb.attributes import Attribute, AttributeContainer, MapAttribute, ListAttribute, UnicodeSetAttribute from pynamodb.connection.base import MetaTable from pynamodb.connection.table import TableConnection from pynamodb.connection.util import pythonic @@ -233,6 +233,95 @@ def __init__(self, hash_key=None, range_key=None, **attrs): attrs[self._dynamo_to_python_attr(range_keyname)] = range_key self._set_attributes(**attrs) + @classmethod + def fix_unicode_set_attributes(cls, + get_save_kwargs, + read_capacity_to_consume_per_second=10, + max_sleep_between_retry=10, + max_consecutive_exceptions=30): + """ + This function performs a rate limited scan of the table and re-serializes any UnicodeSetAttributes. + + See https://github.com/pynamodb/PynamoDB/issues/377 for why this is necessary. + + :param get_save_kwargs: A callback function that is passed a model and should return the kwargs + used when conditionally saving the item + :param read_capacity_to_consume_per_second: Amount of read capacity to consume + every second + :param max_sleep_between_retry: Max value for sleep in seconds in between scans during + throttling/rate limit scenarios + :param max_consecutive_exceptions: Max number of consecutive provision throughput exceeded + exceptions for scan to exit + """ + + if not cls._has_unicode_set_attribute(): + return + + items = cls.rate_limited_scan( + read_capacity_to_consume_per_second=read_capacity_to_consume_per_second, + max_sleep_between_retry=max_sleep_between_retry, + max_consecutive_exceptions=max_consecutive_exceptions + ) + for item in items: + save_kwargs = get_save_kwargs(item) + item.save(**save_kwargs) + + @classmethod + def _has_unicode_set_attribute(cls): + for attr_value in cls._get_attributes().values(): + if isinstance(attr_value, UnicodeSetAttribute): + return True + if isinstance(attr_value, MapAttribute): + return attr_value._has_unicode_set_attribute() + return False + + @classmethod + def needs_unicode_set_fix(cls, + read_capacity_to_consume_per_second=10, + max_sleep_between_retry=10, + max_consecutive_exceptions=30): + + if not cls._has_unicode_set_attribute(): + return False + + scan_result = cls._get_connection().rate_limited_scan( + read_capacity_to_consume_per_second=read_capacity_to_consume_per_second, + max_sleep_between_retry=max_sleep_between_retry, + max_consecutive_exceptions=max_consecutive_exceptions, + ) + + ret_val = False + for item in scan_result: + ret_val |= cls._has_json_unicode_set_value(item) + return ret_val + + @classmethod + def _has_json_unicode_set_value(cls, data): + ret_val = False + for name, attr in data.items(): + # attr should be a map of attribute type to value + (attr_type, attr_value), = attr.items() + if attr_type == 'M': + ret_val |= cls._has_json_unicode_set_value(attr_value) + elif attr_type == 'L': + for value in attr_value: + (at, av), = value.items() + ret_val |= cls._attr_has_json_unicode_set_value(at, av) + else: + ret_val |= cls._attr_has_json_unicode_set_value(attr_type, attr_value) + return ret_val + + @classmethod + def _attr_has_json_unicode_set_value(cls, attr_type, value): + if attr_type != 'SS': + return False + for val in value: + try: + result = json.loads(val) + except ValueError: + return False + return True + @classmethod def has_map_or_list_attributes(cls): for attr_value in cls._get_attributes().values(): diff --git a/pynamodb/tests/test_model.py b/pynamodb/tests/test_model.py index 8fb2c85d1..ab1256032 100644 --- a/pynamodb/tests/test_model.py +++ b/pynamodb/tests/test_model.py @@ -3619,3 +3619,40 @@ def test_subclassed_map_attribute_with_map_attribute_member_with_initialized_ins self.assertEquals(actual.left.left.value, left_instance.left.value) self.assertEquals(actual.right.right.left.value, right_instance.right.left.value) self.assertEquals(actual.right.right.value, right_instance.right.value) + + +class JSONUnicodeSetTestCase(TestCase): + + def test_needs_unidode_set_fix_false(self): + test_item = { + 'string_set_attr': { + 'SS': ['a', 'b'] + }, + 'map_attr': {'M': { + 'foo': {'S': 'bar'}, + 'num': {'N': '1'}, + 'bool_type': {'BOOL': True}, + 'other_b_type': {'BOOL': False}, + 'floaty': {'N': '1.2'}, + 'listy': {'L': [{'N': '1'}, {'N': '2'}, {'N': '12345678909876543211234234324234'}]}, + 'mapy': {'M': {'baz': {'S': 'bongo'}}} + }} + } + assert Model._has_json_unicode_set_value(test_item) == False + + def test_needs_unidode_set_fix_true(self): + test_item = { + 'string_set_attr': { + 'SS': ['a', 'b'] + }, + 'map_attr': {'M': { + 'foo': {'S': 'bar'}, + 'num': {'N': '1'}, + 'bool_type': {'BOOL': True}, + 'other_b_type': {'BOOL': False}, + 'floaty': {'N': '1.2'}, + 'listy': {'L': [{'N': '1'}, {'N': '2'}, {'SS': ['"a"', '"b"']}]}, + 'mapy': {'M': {'baz': {'S': 'bongo'}}} + }} + } + assert Model._has_json_unicode_set_value(test_item) == True