diff --git a/google/cloud/ndb/model.py b/google/cloud/ndb/model.py index a64dfa04..447176c3 100644 --- a/google/cloud/ndb/model.py +++ b/google/cloud/ndb/model.py @@ -251,6 +251,7 @@ class Person(Model): import inspect import json import pickle +import six import zlib from google.cloud.datastore import entity as ds_entity_module @@ -2188,15 +2189,13 @@ def _validate(self, value): return float(value) -class _CompressedValue: +class _CompressedValue(six.binary_type): """A marker object wrapping compressed values. Args: z_val (bytes): A return value of ``zlib.compress``. """ - __slots__ = ("z_val",) - def __init__(self, z_val): self.z_val = z_val @@ -2355,6 +2354,9 @@ def _from_base_type(self, value): indicate that the value didn't need to be unwrapped and decompressed. """ + if self._compressed and not isinstance(value, _CompressedValue): + value = _CompressedValue(value) + if isinstance(value, _CompressedValue): return zlib.decompress(value.z_val) diff --git a/tests/system/test_crud.py b/tests/system/test_crud.py index 23bb5c12..5241ba76 100644 --- a/tests/system/test_crud.py +++ b/tests/system/test_crud.py @@ -285,6 +285,36 @@ class SomeKind(ndb.Model): dispose_of(key._key) +@pytest.mark.usefixtures("client_context") +def test_compressed_json_property(dispose_of, ds_client): + class SomeKind(ndb.Model): + foo = ndb.JsonProperty(compressed=True) + + foo = {str(i): i for i in range(500)} + entity = SomeKind(foo=foo) + key = entity.put() + + retrieved = key.get() + assert retrieved.foo == foo + + dispose_of(key._key) + + +@pytest.mark.usefixtures("client_context") +def test_compressed_blob_property(dispose_of, ds_client): + class SomeKind(ndb.Model): + foo = ndb.BlobProperty(compressed=True) + + foo = b"abc" * 100 + entity = SomeKind(foo=foo) + key = entity.put() + + retrieved = key.get() + assert retrieved.foo == foo + + dispose_of(key._key) + + @pytest.mark.usefixtures("client_context") def test_large_pickle_property(dispose_of, ds_client): class SomeKind(ndb.Model): diff --git a/tests/unit/test_model.py b/tests/unit/test_model.py index 183f3702..9b3ec09f 100644 --- a/tests/unit/test_model.py +++ b/tests/unit/test_model.py @@ -1677,6 +1677,16 @@ def test__from_base_type(): assert converted == original + @staticmethod + def test__from_base_type_no_compressed_value(): + prop = model.BlobProperty(name="blob") + original = b"abc" * 10 + value = zlib.compress(original) + prop._compressed = True + converted = prop._from_base_type(value) + + assert converted == original + @staticmethod def test__from_base_type_no_convert(): prop = model.BlobProperty(name="blob")