Skip to content

Commit 014b016

Browse files
committed
Merge pull request #232 from tseaver/support-entity-protobufs
Support entity protobufs
2 parents 40c1852 + fef83a1 commit 014b016

File tree

7 files changed

+236
-22
lines changed

7 files changed

+236
-22
lines changed
Lines changed: 51 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,39 @@
1-
"""Helper methods for dealing with Cloud Datastore's Protobuf API."""
1+
"""Helper functions for dealing with Cloud Datastore's Protobuf API.
2+
3+
These functions are *not* part of the API.
4+
"""
25
import calendar
36
from datetime import datetime, timedelta
47

58
from google.protobuf.internal.type_checkers import Int64ValueChecker
69
import pytz
710

11+
from gcloud.datastore.entity import Entity
812
from gcloud.datastore.key import Key
913

1014
INT_VALUE_CHECKER = Int64ValueChecker()
1115

1216

13-
def get_protobuf_attribute_and_value(val):
17+
def _get_protobuf_attribute_and_value(val):
1418
"""Given a value, return the protobuf attribute name and proper value.
1519
1620
The Protobuf API uses different attribute names
1721
based on value types rather than inferring the type.
18-
This method simply determines the proper attribute name
22+
This function simply determines the proper attribute name
1923
based on the type of the value provided
2024
and returns the attribute name
2125
as well as a properly formatted value.
2226
2327
Certain value types need to be coerced into a different type (such as a
2428
`datetime.datetime` into an integer timestamp, or a
2529
`gcloud.datastore.key.Key` into a Protobuf representation.
26-
This method handles that for you.
30+
This function handles that for you.
2731
2832
For example:
2933
30-
>>> get_protobuf_attribute_and_value(1234)
34+
>>> _get_protobuf_attribute_and_value(1234)
3135
('integer_value', 1234)
32-
>>> get_protobuf_attribute_and_value('my_string')
36+
>>> _get_protobuf_attribute_and_value('my_string')
3337
('string_value', 'my_string')
3438
3539
:type val: `datetime.datetime`, :class:`gcloud.datastore.key.Key`,
@@ -60,18 +64,20 @@ def get_protobuf_attribute_and_value(val):
6064
name, value = 'integer', long(val) # Always cast to a long.
6165
elif isinstance(val, basestring):
6266
name, value = 'string', val
67+
elif isinstance(val, Entity):
68+
name, value = 'entity', val
6369
else:
6470
raise ValueError("Unknown protobuf attr type %s" % type(val))
6571

6672
return name + '_value', value
6773

6874

69-
def get_value_from_protobuf(pb):
75+
def _get_value_from_protobuf(pb):
7076
"""Given a protobuf for a Property, get the correct value.
7177
7278
The Cloud Datastore Protobuf API returns a Property Protobuf
7379
which has one value set and the rest blank.
74-
This method retrieves the the one value provided.
80+
This function retrieves the the one value provided.
7581
7682
Some work is done to coerce the return value into a more useful type
7783
(particularly in the case of a timestamp value, or a key value).
@@ -103,5 +109,42 @@ def get_value_from_protobuf(pb):
103109
elif pb.value.HasField('string_value'):
104110
return pb.value.string_value
105111

112+
elif pb.value.HasField('entity_value'):
113+
return Entity.from_protobuf(pb.value.entity_value)
114+
106115
else:
107116
return None
117+
118+
119+
def _set_protobuf_value(value_pb, val):
120+
"""Assign 'val' to the correct subfield of 'value_pb'.
121+
122+
The Protobuf API uses different attribute names
123+
based on value types rather than inferring the type.
124+
125+
Some value types (entities, keys, lists) cannot be directly assigned;
126+
this function handles them correctly.
127+
128+
:type value_pb: :class:`gcloud.datastore.datastore_v1_pb2.Value`
129+
:param value_pb: The value protobuf to which the value is being assigned.
130+
131+
:type val: `datetime.datetime`, bool, float, integer, string
132+
:class:`gcloud.datastore.key.Key`,
133+
:class:`gcloud.datastore.entity.Entity`,
134+
:param val: The value to be assigned.
135+
"""
136+
attr, val = _get_protobuf_attribute_and_value(val)
137+
if attr == 'key_value':
138+
value_pb.key_value.CopyFrom(val)
139+
elif attr == 'entity_value':
140+
e_pb = value_pb.entity_value
141+
e_pb.Clear()
142+
key = val.key()
143+
if key is not None:
144+
e_pb.key.CopyFrom(key.to_protobuf())
145+
for k, v in val.items():
146+
p_pb = e_pb.property.add()
147+
p_pb.name = k
148+
_set_protobuf_value(p_pb.value, v)
149+
else: # scalar, just assign
150+
setattr(value_pb, attr, val)

gcloud/datastore/connection.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from gcloud import connection
22
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
3-
from gcloud.datastore import helpers
3+
from gcloud.datastore import _helpers
44
from gcloud.datastore.dataset import Dataset
55

66

@@ -323,8 +323,7 @@ def save_entity(self, dataset_id, key_pb, properties):
323323
prop.name = name
324324

325325
# Set the appropriate value.
326-
pb_attr, pb_value = helpers.get_protobuf_attribute_and_value(value)
327-
setattr(prop.value, pb_attr, pb_value)
326+
_helpers._set_protobuf_value(prop.value, value)
328327

329328
# If this is in a transaction, we should just return True. The
330329
# transaction will handle assigning any keys as necessary.

gcloud/datastore/entity.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -151,13 +151,13 @@ def from_protobuf(cls, pb, dataset=None): # pylint: disable=invalid-name
151151
"""
152152

153153
# This is here to avoid circular imports.
154-
from gcloud.datastore import helpers
154+
from gcloud.datastore import _helpers
155155

156156
key = Key.from_protobuf(pb.key, dataset=dataset)
157157
entity = cls.from_key(key)
158158

159159
for property_pb in pb.property:
160-
value = helpers.get_value_from_protobuf(property_pb)
160+
value = _helpers._get_value_from_protobuf(property_pb)
161161
entity[property_pb.name] = value
162162

163163
return entity

gcloud/datastore/query.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
import copy
33

44
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
5-
from gcloud.datastore import helpers
5+
from gcloud.datastore import _helpers
66
from gcloud.datastore.entity import Entity
77
from gcloud.datastore.key import Key
88

@@ -131,8 +131,7 @@ def filter(self, expression, value):
131131
property_filter.operator = operator
132132

133133
# Set the value to filter on based on the type.
134-
attr_name, pb_value = helpers.get_protobuf_attribute_and_value(value)
135-
setattr(property_filter.value, attr_name, pb_value)
134+
_helpers._set_protobuf_value(property_filter.value, value)
136135
return clone
137136

138137
def ancestor(self, ancestor):
Lines changed: 130 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import unittest2
22

33

4-
class Test_get_protobuf_attribute_and_value(unittest2.TestCase):
4+
class Test__get_protobuf_attribute_and_value(unittest2.TestCase):
55

66
def _callFUT(self, val):
7-
from gcloud.datastore.helpers import get_protobuf_attribute_and_value
7+
from gcloud.datastore._helpers import _get_protobuf_attribute_and_value
88

9-
return get_protobuf_attribute_and_value(val)
9+
return _get_protobuf_attribute_and_value(val)
1010

1111
def test_datetime_naive(self):
1212
import calendar
@@ -83,16 +83,23 @@ def test_unicode(self):
8383
self.assertEqual(name, 'string_value')
8484
self.assertEqual(value, u'str')
8585

86+
def test_entity(self):
87+
from gcloud.datastore.entity import Entity
88+
entity = Entity()
89+
name, value = self._callFUT(entity)
90+
self.assertEqual(name, 'entity_value')
91+
self.assertTrue(value is entity)
92+
8693
def test_object(self):
8794
self.assertRaises(ValueError, self._callFUT, object())
8895

8996

90-
class Test_get_value_from_protobuf(unittest2.TestCase):
97+
class Test__get_value_from_protobuf(unittest2.TestCase):
9198

9299
def _callFUT(self, pb):
93-
from gcloud.datastore.helpers import get_value_from_protobuf
100+
from gcloud.datastore._helpers import _get_value_from_protobuf
94101

95-
return get_value_from_protobuf(pb)
102+
return _get_value_from_protobuf(pb)
96103

97104
def _makePB(self, attr_name, value):
98105
from gcloud.datastore.datastore_v1_pb2 import Property
@@ -146,7 +153,124 @@ def test_unicode(self):
146153
pb = self._makePB('string_value', u'str')
147154
self.assertEqual(self._callFUT(pb), u'str')
148155

156+
def test_entity(self):
157+
from gcloud.datastore.datastore_v1_pb2 import Property
158+
from gcloud.datastore.entity import Entity
159+
160+
pb = Property()
161+
entity_pb = pb.value.entity_value
162+
prop_pb = entity_pb.property.add()
163+
prop_pb.name = 'foo'
164+
prop_pb.value.string_value = 'Foo'
165+
entity = self._callFUT(pb)
166+
self.assertTrue(isinstance(entity, Entity))
167+
self.assertEqual(entity['foo'], 'Foo')
168+
149169
def test_unknown(self):
150170
from gcloud.datastore.datastore_v1_pb2 import Property
171+
151172
pb = Property()
152173
self.assertEqual(self._callFUT(pb), None) # XXX desirable?
174+
175+
176+
class Test_set_protobuf_value(unittest2.TestCase):
177+
178+
def _callFUT(self, value_pb, val):
179+
from gcloud.datastore._helpers import _set_protobuf_value
180+
181+
return _set_protobuf_value(value_pb, val)
182+
183+
def _makePB(self):
184+
from gcloud.datastore.datastore_v1_pb2 import Value
185+
186+
return Value()
187+
188+
def test_datetime(self):
189+
import calendar
190+
import datetime
191+
import pytz
192+
193+
pb = self._makePB()
194+
utc = datetime.datetime(2014, 9, 16, 10, 19, 32, 4375, pytz.utc)
195+
self._callFUT(pb, utc)
196+
value = pb.timestamp_microseconds_value
197+
self.assertEqual(value / 1000000, calendar.timegm(utc.timetuple()))
198+
self.assertEqual(value % 1000000, 4375)
199+
200+
def test_key(self):
201+
from gcloud.datastore.dataset import Dataset
202+
from gcloud.datastore.key import Key
203+
204+
_DATASET = 'DATASET'
205+
_KIND = 'KIND'
206+
_ID = 1234
207+
_PATH = [{'kind': _KIND, 'id': _ID}]
208+
pb = self._makePB()
209+
key = Key(dataset=Dataset(_DATASET), path=_PATH)
210+
self._callFUT(pb, key)
211+
value = pb.key_value
212+
self.assertEqual(value, key.to_protobuf())
213+
214+
def test_bool(self):
215+
pb = self._makePB()
216+
self._callFUT(pb, False)
217+
value = pb.boolean_value
218+
self.assertEqual(value, False)
219+
220+
def test_float(self):
221+
pb = self._makePB()
222+
self._callFUT(pb, 3.1415926)
223+
value = pb.double_value
224+
self.assertEqual(value, 3.1415926)
225+
226+
def test_int(self):
227+
pb = self._makePB()
228+
self._callFUT(pb, 42)
229+
value = pb.integer_value
230+
self.assertEqual(value, 42)
231+
232+
def test_long(self):
233+
pb = self._makePB()
234+
must_be_long = (1 << 63) - 1
235+
self._callFUT(pb, must_be_long)
236+
value = pb.integer_value
237+
self.assertEqual(value, must_be_long)
238+
239+
def test_native_str(self):
240+
pb = self._makePB()
241+
self._callFUT(pb, 'str')
242+
value = pb.string_value
243+
self.assertEqual(value, 'str')
244+
245+
def test_unicode(self):
246+
pb = self._makePB()
247+
self._callFUT(pb, u'str')
248+
value = pb.string_value
249+
self.assertEqual(value, u'str')
250+
251+
def test_entity_empty_wo_key(self):
252+
from gcloud.datastore.entity import Entity
253+
254+
pb = self._makePB()
255+
entity = Entity()
256+
self._callFUT(pb, entity)
257+
value = pb.entity_value
258+
self.assertEqual(value.key.SerializeToString(), '')
259+
props = list(value.property)
260+
self.assertEqual(len(props), 0)
261+
262+
def test_entity_w_key(self):
263+
from gcloud.datastore.entity import Entity
264+
from gcloud.datastore.key import Key
265+
266+
pb = self._makePB()
267+
key = Key(path=[{'kind': 'KIND', 'id': 123}])
268+
entity = Entity().key(key)
269+
entity['foo'] = 'Foo'
270+
self._callFUT(pb, entity)
271+
value = pb.entity_value
272+
self.assertEqual(value.key, key.to_protobuf())
273+
props = list(value.property)
274+
self.assertEqual(len(props), 1)
275+
self.assertEqual(props[0].name, 'foo')
276+
self.assertEqual(props[0].value.string_value, 'Foo')

gcloud/datastore/test_connection.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -699,6 +699,32 @@ def mutation(self):
699699
mutation = conn.mutation()
700700
self.assertEqual(len(mutation.upsert), 1)
701701

702+
def test_save_entity_w_transaction_nested_entity(self):
703+
from gcloud.datastore.connection import datastore_pb
704+
from gcloud.datastore.dataset import Dataset
705+
from gcloud.datastore.entity import Entity
706+
from gcloud.datastore.key import Key
707+
708+
mutation = datastore_pb.Mutation()
709+
710+
class Xact(object):
711+
def mutation(self):
712+
return mutation
713+
DATASET_ID = 'DATASET'
714+
nested = Entity()
715+
nested['bar'] = 'Bar'
716+
key_pb = Key(dataset=Dataset(DATASET_ID),
717+
path=[{'kind': 'Kind', 'id': 1234}]).to_protobuf()
718+
rsp_pb = datastore_pb.CommitResponse()
719+
conn = self._makeOne()
720+
conn.transaction(Xact())
721+
http = conn._http = Http({'status': '200'}, rsp_pb.SerializeToString())
722+
result = conn.save_entity(DATASET_ID, key_pb, {'foo': nested})
723+
self.assertEqual(result, True)
724+
self.assertEqual(http._called_with, None)
725+
mutation = conn.mutation()
726+
self.assertEqual(len(mutation.upsert), 1)
727+
702728
def test_delete_entities_wo_transaction(self):
703729
from gcloud.datastore.connection import datastore_pb
704730
from gcloud.datastore.dataset import Dataset

gcloud/datastore/test_query.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,29 @@ def test_filter_w_known_operator(self):
7171
self.assertEqual(p_pb.property.name, 'firstname')
7272
self.assertEqual(p_pb.value.string_value, 'John')
7373

74+
def test_filter_w_known_operator_and_entity(self):
75+
import operator
76+
from gcloud.datastore.entity import Entity
77+
query = self._makeOne()
78+
other = Entity()
79+
other['firstname'] = 'John'
80+
other['lastname'] = 'Smith'
81+
after = query.filter('other =', other)
82+
self.assertFalse(after is query)
83+
self.assertTrue(isinstance(after, self._getTargetClass()))
84+
q_pb = after.to_protobuf()
85+
self.assertEqual(q_pb.filter.composite_filter.operator, 1) # AND
86+
f_pb, = list(q_pb.filter.composite_filter.filter)
87+
p_pb = f_pb.property_filter
88+
self.assertEqual(p_pb.property.name, 'other')
89+
other_pb = p_pb.value.entity_value
90+
props = sorted(other_pb.property, key=operator.attrgetter('name'))
91+
self.assertEqual(len(props), 2)
92+
self.assertEqual(props[0].name, 'firstname')
93+
self.assertEqual(props[0].value.string_value, 'John')
94+
self.assertEqual(props[1].name, 'lastname')
95+
self.assertEqual(props[1].value.string_value, 'Smith')
96+
7497
def test_ancestor_w_non_key_non_list(self):
7598
query = self._makeOne()
7699
# XXX s.b. ValueError

0 commit comments

Comments
 (0)