Support entity protobufs

gcloud/datastore/connection.py

@@ -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/helpers.py

@@ -5,6 +5,7 @@
 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()
@@ -60,6 +61,8 @@ 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))
 
@@ -103,5 +106,44 @@ 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)  # XXX dataset?
+
     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: 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 None:
+            e_pb.key.CopyFrom(Key().to_protobuf())
+        else:
+            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/query.py

@@ -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_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_helpers.py

@@ -83,6 +83,13 @@ 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())
 
@@ -146,7 +153,125 @@ 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
+        from gcloud.datastore.key import Key
+
+        pb = self._makePB()
+        entity = Entity()
+        self._callFUT(pb, entity)
+        value = pb.entity_value
+        self.assertEqual(value.key, Key().to_protobuf())
+        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_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