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