Merge pull request #232 from tseaver/support-entity-protobufs

Support entity protobufs

gcloud/datastore/helpers.py → gcloud/datastore/_helpers.py

@@ -1,35 +1,39 @@
-"""Helper methods for dealing with Cloud Datastore's Protobuf API."""
+"""Helper functions for dealing with Cloud Datastore's Protobuf API.
+
+These functions are *not* part of the API.
+"""
 import calendar
 from datetime import datetime, timedelta
 
 from google.protobuf.internal.type_checkers import Int64ValueChecker
 import pytz
 
+from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 
 INT_VALUE_CHECKER = Int64ValueChecker()
 
 
-def get_protobuf_attribute_and_value(val):
+def _get_protobuf_attribute_and_value(val):
     """Given a value, return the protobuf attribute name and proper value.
 
     The Protobuf API uses different attribute names
     based on value types rather than inferring the type.
-    This method simply determines the proper attribute name
+    This function simply determines the proper attribute name
     based on the type of the value provided
     and returns the attribute name
     as well as a properly formatted value.
 
     Certain value types need to be coerced into a different type (such as a
     `datetime.datetime` into an integer timestamp, or a
     `gcloud.datastore.key.Key` into a Protobuf representation.
-    This method handles that for you.
+    This function handles that for you.
 
     For example:
 
-    >>> get_protobuf_attribute_and_value(1234)
+    >>> _get_protobuf_attribute_and_value(1234)
     ('integer_value', 1234)
-    >>> get_protobuf_attribute_and_value('my_string')
+    >>> _get_protobuf_attribute_and_value('my_string')
     ('string_value', 'my_string')
 
     :type val: `datetime.datetime`, :class:`gcloud.datastore.key.Key`,
@@ -60,18 +64,20 @@ def get_protobuf_attribute_and_value(val):
         name, value = 'integer', long(val)  # Always cast to a long.
     elif isinstance(val, basestring):
         name, value = 'string', val
+    elif isinstance(val, Entity):
+        name, value = 'entity', val
     else:
         raise ValueError("Unknown protobuf attr type %s" % type(val))
 
     return name + '_value', value
 
 
-def get_value_from_protobuf(pb):
+def _get_value_from_protobuf(pb):
     """Given a protobuf for a Property, get the correct value.
 
     The Cloud Datastore Protobuf API returns a Property Protobuf
     which has one value set and the rest blank.
-    This method retrieves the the one value provided.
+    This function retrieves the the one value provided.
 
     Some work is done to coerce the return value into a more useful type
     (particularly in the case of a timestamp value, or a key value).
@@ -103,5 +109,42 @@ def get_value_from_protobuf(pb):
     elif pb.value.HasField('string_value'):
         return pb.value.string_value
 
+    elif pb.value.HasField('entity_value'):
+        return Entity.from_protobuf(pb.value.entity_value)
+
     else:
         return None
+
+
+def _set_protobuf_value(value_pb, val):
+    """Assign 'val' to the correct subfield of 'value_pb'.
+
+    The Protobuf API uses different attribute names
+    based on value types rather than inferring the type.
+
+    Some value types (entities, keys, lists) cannot be directly assigned;
+    this function handles them correctly.
+
+    :type value_pb: :class:`gcloud.datastore.datastore_v1_pb2.Value`
+    :param value_pb: The value protobuf to which the value is being assigned.
+
+    :type val: `datetime.datetime`, bool, float, integer, string
+               :class:`gcloud.datastore.key.Key`,
+               :class:`gcloud.datastore.entity.Entity`,
+    :param val: The value to be assigned.
+    """
+    attr, val = _get_protobuf_attribute_and_value(val)
+    if attr == 'key_value':
+        value_pb.key_value.CopyFrom(val)
+    elif attr == 'entity_value':
+        e_pb = value_pb.entity_value
+        e_pb.Clear()
+        key = val.key()
+        if key is not None:
+            e_pb.key.CopyFrom(key.to_protobuf())
+        for k, v in val.items():
+            p_pb = e_pb.property.add()
+            p_pb.name = k
+            _set_protobuf_value(p_pb.value, v)
+    else:  # scalar, just assign
+        setattr(value_pb, attr, val)

gcloud/datastore/connection.py

@@ -1,6 +1,6 @@
 from gcloud import connection
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
-from gcloud.datastore import helpers
+from gcloud.datastore import _helpers
 from gcloud.datastore.dataset import Dataset
 
 
@@ -323,8 +323,7 @@ def save_entity(self, dataset_id, key_pb, properties):
             prop.name = name
 
             # Set the appropriate value.
-            pb_attr, pb_value = helpers.get_protobuf_attribute_and_value(value)
-            setattr(prop.value, pb_attr, pb_value)
+            _helpers._set_protobuf_value(prop.value, value)
 
         # If this is in a transaction, we should just return True. The
         # transaction will handle assigning any keys as necessary.

gcloud/datastore/entity.py

@@ -151,13 +151,13 @@ def from_protobuf(cls, pb, dataset=None):  # pylint: disable=invalid-name
         """
 
         # This is here to avoid circular imports.
-        from gcloud.datastore import helpers
+        from gcloud.datastore import _helpers
 
         key = Key.from_protobuf(pb.key, dataset=dataset)
         entity = cls.from_key(key)
 
         for property_pb in pb.property:
-            value = helpers.get_value_from_protobuf(property_pb)
+            value = _helpers._get_value_from_protobuf(property_pb)
             entity[property_pb.name] = value
 
         return entity

gcloud/datastore/query.py

@@ -2,7 +2,7 @@
 import copy
 
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
-from gcloud.datastore import helpers
+from gcloud.datastore import _helpers
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.key import Key
 
@@ -131,8 +131,7 @@ def filter(self, expression, value):
         property_filter.operator = operator
 
         # Set the value to filter on based on the type.
-        attr_name, pb_value = helpers.get_protobuf_attribute_and_value(value)
-        setattr(property_filter.value, attr_name, pb_value)
+        _helpers._set_protobuf_value(property_filter.value, value)
         return clone
 
     def ancestor(self, ancestor):

gcloud/datastore/test_helpers.py → gcloud/datastore/test__helpers.py

@@ -1,12 +1,12 @@
 import unittest2
 
 
-class Test_get_protobuf_attribute_and_value(unittest2.TestCase):
+class Test__get_protobuf_attribute_and_value(unittest2.TestCase):
 
     def _callFUT(self, val):
-        from gcloud.datastore.helpers import get_protobuf_attribute_and_value
+        from gcloud.datastore._helpers import _get_protobuf_attribute_and_value
 
-        return get_protobuf_attribute_and_value(val)
+        return _get_protobuf_attribute_and_value(val)
 
     def test_datetime_naive(self):
         import calendar
@@ -83,16 +83,23 @@ def test_unicode(self):
         self.assertEqual(name, 'string_value')
         self.assertEqual(value, u'str')
 
+    def test_entity(self):
+        from gcloud.datastore.entity import Entity
+        entity = Entity()
+        name, value = self._callFUT(entity)
+        self.assertEqual(name, 'entity_value')
+        self.assertTrue(value is entity)
+
     def test_object(self):
         self.assertRaises(ValueError, self._callFUT, object())
 
 
-class Test_get_value_from_protobuf(unittest2.TestCase):
+class Test__get_value_from_protobuf(unittest2.TestCase):
 
     def _callFUT(self, pb):
-        from gcloud.datastore.helpers import get_value_from_protobuf
+        from gcloud.datastore._helpers import _get_value_from_protobuf
 
-        return get_value_from_protobuf(pb)
+        return _get_value_from_protobuf(pb)
 
     def _makePB(self, attr_name, value):
         from gcloud.datastore.datastore_v1_pb2 import Property
@@ -146,7 +153,124 @@ def test_unicode(self):
         pb = self._makePB('string_value', u'str')
         self.assertEqual(self._callFUT(pb), u'str')
 
+    def test_entity(self):
+        from gcloud.datastore.datastore_v1_pb2 import Property
+        from gcloud.datastore.entity import Entity
+
+        pb = Property()
+        entity_pb = pb.value.entity_value
+        prop_pb = entity_pb.property.add()
+        prop_pb.name = 'foo'
+        prop_pb.value.string_value = 'Foo'
+        entity = self._callFUT(pb)
+        self.assertTrue(isinstance(entity, Entity))
+        self.assertEqual(entity['foo'], 'Foo')
+
     def test_unknown(self):
         from gcloud.datastore.datastore_v1_pb2 import Property
+
         pb = Property()
         self.assertEqual(self._callFUT(pb), None)  # XXX desirable?
+
+
+class Test_set_protobuf_value(unittest2.TestCase):
+
+    def _callFUT(self, value_pb, val):
+        from gcloud.datastore._helpers import _set_protobuf_value
+
+        return _set_protobuf_value(value_pb, val)
+
+    def _makePB(self):
+        from gcloud.datastore.datastore_v1_pb2 import Value
+
+        return Value()
+
+    def test_datetime(self):
+        import calendar
+        import datetime
+        import pytz
+
+        pb = self._makePB()
+        utc = datetime.datetime(2014, 9, 16, 10, 19, 32, 4375, pytz.utc)
+        self._callFUT(pb, utc)
+        value = pb.timestamp_microseconds_value
+        self.assertEqual(value / 1000000, calendar.timegm(utc.timetuple()))
+        self.assertEqual(value % 1000000, 4375)
+
+    def test_key(self):
+        from gcloud.datastore.dataset import Dataset
+        from gcloud.datastore.key import Key
+
+        _DATASET = 'DATASET'
+        _KIND = 'KIND'
+        _ID = 1234
+        _PATH = [{'kind': _KIND, 'id': _ID}]
+        pb = self._makePB()
+        key = Key(dataset=Dataset(_DATASET), path=_PATH)
+        self._callFUT(pb, key)
+        value = pb.key_value
+        self.assertEqual(value, key.to_protobuf())
+
+    def test_bool(self):
+        pb = self._makePB()
+        self._callFUT(pb, False)
+        value = pb.boolean_value
+        self.assertEqual(value, False)
+
+    def test_float(self):
+        pb = self._makePB()
+        self._callFUT(pb, 3.1415926)
+        value = pb.double_value
+        self.assertEqual(value, 3.1415926)
+
+    def test_int(self):
+        pb = self._makePB()
+        self._callFUT(pb, 42)
+        value = pb.integer_value
+        self.assertEqual(value, 42)
+
+    def test_long(self):
+        pb = self._makePB()
+        must_be_long = (1 << 63) - 1
+        self._callFUT(pb, must_be_long)
+        value = pb.integer_value
+        self.assertEqual(value, must_be_long)
+
+    def test_native_str(self):
+        pb = self._makePB()
+        self._callFUT(pb, 'str')
+        value = pb.string_value
+        self.assertEqual(value, 'str')
+
+    def test_unicode(self):
+        pb = self._makePB()
+        self._callFUT(pb, u'str')
+        value = pb.string_value
+        self.assertEqual(value, u'str')
+
+    def test_entity_empty_wo_key(self):
+        from gcloud.datastore.entity import Entity
+
+        pb = self._makePB()
+        entity = Entity()
+        self._callFUT(pb, entity)
+        value = pb.entity_value
+        self.assertEqual(value.key.SerializeToString(), '')
+        props = list(value.property)
+        self.assertEqual(len(props), 0)
+
+    def test_entity_w_key(self):
+        from gcloud.datastore.entity import Entity
+        from gcloud.datastore.key import Key
+
+        pb = self._makePB()
+        key = Key(path=[{'kind': 'KIND', 'id': 123}])
+        entity = Entity().key(key)
+        entity['foo'] = 'Foo'
+        self._callFUT(pb, entity)
+        value = pb.entity_value
+        self.assertEqual(value.key, key.to_protobuf())
+        props = list(value.property)
+        self.assertEqual(len(props), 1)
+        self.assertEqual(props[0].name, 'foo')
+        self.assertEqual(props[0].value.string_value, 'Foo')

gcloud/datastore/test_connection.py

@@ -699,6 +699,32 @@ def mutation(self):
         mutation = conn.mutation()
         self.assertEqual(len(mutation.upsert), 1)
 
+    def test_save_entity_w_transaction_nested_entity(self):
+        from gcloud.datastore.connection import datastore_pb
+        from gcloud.datastore.dataset import Dataset
+        from gcloud.datastore.entity import Entity
+        from gcloud.datastore.key import Key
+
+        mutation = datastore_pb.Mutation()
+
+        class Xact(object):
+            def mutation(self):
+                return mutation
+        DATASET_ID = 'DATASET'
+        nested = Entity()
+        nested['bar'] = 'Bar'
+        key_pb = Key(dataset=Dataset(DATASET_ID),
+                     path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
+        rsp_pb = datastore_pb.CommitResponse()
+        conn = self._makeOne()
+        conn.transaction(Xact())
+        http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
+        result = conn.save_entity(DATASET_ID, key_pb, {'foo': nested})
+        self.assertEqual(result, True)
+        self.assertEqual(http._called_with, None)
+        mutation = conn.mutation()
+        self.assertEqual(len(mutation.upsert), 1)
+
     def test_delete_entities_wo_transaction(self):
         from gcloud.datastore.connection import datastore_pb
         from gcloud.datastore.dataset import Dataset

gcloud/datastore/test_query.py

@@ -71,6 +71,29 @@ def test_filter_w_known_operator(self):
         self.assertEqual(p_pb.property.name, 'firstname')
         self.assertEqual(p_pb.value.string_value, 'John')
 
+    def test_filter_w_known_operator_and_entity(self):
+        import operator
+        from gcloud.datastore.entity import Entity
+        query = self._makeOne()
+        other = Entity()
+        other['firstname'] = 'John'
+        other['lastname'] = 'Smith'
+        after = query.filter('other =', other)
+        self.assertFalse(after is query)
+        self.assertTrue(isinstance(after, self._getTargetClass()))
+        q_pb = after.to_protobuf()
+        self.assertEqual(q_pb.filter.composite_filter.operator, 1)  # AND
+        f_pb, = list(q_pb.filter.composite_filter.filter)
+        p_pb = f_pb.property_filter
+        self.assertEqual(p_pb.property.name, 'other')
+        other_pb = p_pb.value.entity_value
+        props = sorted(other_pb.property, key=operator.attrgetter('name'))
+        self.assertEqual(len(props), 2)
+        self.assertEqual(props[0].name, 'firstname')
+        self.assertEqual(props[0].value.string_value, 'John')
+        self.assertEqual(props[1].name, 'lastname')
+        self.assertEqual(props[1].value.string_value, 'Smith')
+
     def test_ancestor_w_non_key_non_list(self):
         query = self._makeOne()
         # XXX s.b. ValueError