Skip to content

Commit 0194f30

Browse files
committed
Adding a non-protobuf allocate_ids method on Dataset.
Partially address #336.
1 parent c2ba130 commit 0194f30

File tree

4 files changed

+88
-8
lines changed

4 files changed

+88
-8
lines changed

gcloud/datastore/dataset.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,3 +182,29 @@ def get_entities(self, keys, missing=None, deferred=None):
182182
entities.append(helpers.entity_from_protobuf(
183183
entity_pb, dataset=self))
184184
return entities
185+
186+
def allocate_ids(self, incomplete_key, num_ids):
187+
"""Allocates a list of IDs from a partial key.
188+
189+
:type incomplete_key: A :class:`gcloud.datastore.key.Key`
190+
:param incomplete_key: Partial key to use as base for allocated IDs.
191+
192+
:type num_ids: A :class:`int`.
193+
:param num_ids: The number of IDs to allocate.
194+
195+
:rtype: list of :class:`gcloud.datastore.key.Key`
196+
:return: The (complete) keys allocated with `incomplete_key` as root.
197+
:raises: `ValueError` if `incomplete_key` is not a partial key.
198+
"""
199+
if not incomplete_key.is_partial():
200+
raise ValueError(('Key is not partial.', incomplete_key))
201+
202+
incomplete_key_pb = incomplete_key.to_protobuf()
203+
incomplete_key_pbs = [incomplete_key_pb] * num_ids
204+
205+
allocated_key_pbs = self.connection().allocate_ids(
206+
self.id(), incomplete_key_pbs)
207+
allocated_ids = [allocated_key_pb.path_element[-1].id
208+
for allocated_key_pb in allocated_key_pbs]
209+
return [incomplete_key.id(allocated_id)
210+
for allocated_id in allocated_ids]

gcloud/datastore/test_dataset.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,40 @@ def test_get_entities_hit(self):
213213
self.assertEqual(list(result), ['foo'])
214214
self.assertEqual(result['foo'], 'Foo')
215215

216+
def test_allocate_ids(self):
217+
from gcloud.datastore.test_entity import _Key
218+
219+
INCOMPLETE_KEY = _Key()
220+
PROTO_ID = object()
221+
INCOMPLETE_KEY._key = _KeyProto(PROTO_ID)
222+
INCOMPLETE_KEY._partial = True
223+
224+
CONNECTION = _Connection()
225+
NUM_IDS = 2
226+
DATASET_ID = 'foo'
227+
DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
228+
result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
229+
230+
# Check the IDs returned match _PathElementProto.
231+
self.assertEqual(result, range(NUM_IDS))
232+
233+
# Check connection is called correctly.
234+
self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
235+
self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
236+
237+
# Check the IDs passed to Connection.allocate_ids.
238+
key_paths = [key_pb.path_element[-1].id
239+
for key_pb in CONNECTION._called_key_pbs]
240+
self.assertEqual(key_paths, [PROTO_ID] * NUM_IDS)
241+
242+
def test_allocate_ids_with_complete(self):
243+
from gcloud.datastore.test_entity import _Key
244+
245+
COMPLETE_KEY = _Key()
246+
DATASET = self._makeOne(None)
247+
self.assertRaises(ValueError, DATASET.allocate_ids,
248+
COMPLETE_KEY, 2)
249+
216250

217251
class _Connection(object):
218252
_called_with = None
@@ -230,3 +264,21 @@ def lookup(self, **kw):
230264
if deferred is not None:
231265
deferred.extend(self._deferred)
232266
return self._result
267+
268+
def allocate_ids(self, dataset_id, key_pbs):
269+
self._called_dataset_id = dataset_id
270+
self._called_key_pbs = key_pbs
271+
num_pbs = len(key_pbs)
272+
return [_KeyProto(i) for i in xrange(num_pbs)]
273+
274+
275+
class _PathElementProto(object):
276+
277+
def __init__(self, _id):
278+
self.id = _id
279+
280+
281+
class _KeyProto(object):
282+
283+
def __init__(self, id_):
284+
self.path_element = [_PathElementProto(id_)]

gcloud/datastore/test_entity.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,10 @@ class _Key(object):
240240
_partial = False
241241
_path = None
242242

243+
def id(self, id_to_set):
244+
self._called_id = id_to_set
245+
return id_to_set
246+
243247
def to_protobuf(self):
244248
return self._key
245249

regression/datastore.py

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,18 +42,16 @@ class TestDatastoreAllocateIDs(TestDatastore):
4242

4343
def test_allocate_ids(self):
4444
incomplete_key = datastore.key.Key(path=[{'kind': 'Kind'}])
45-
incomplete_key_pb = incomplete_key.to_protobuf()
46-
incomplete_key_pbs = [incomplete_key_pb] * 10
47-
48-
connection = self.dataset.connection()
49-
allocated_key_pbs = connection.allocate_ids(self.dataset.id(),
50-
incomplete_key_pbs)
51-
allocated_keys = [datastore.helpers.key_from_protobuf(key_pb)
52-
for key_pb in allocated_key_pbs]
45+
allocated_keys = self.dataset.allocate_ids(incomplete_key, 10)
5346
self.assertEqual(len(allocated_keys), 10)
47+
48+
unique_ids = set()
5449
for key in allocated_keys:
50+
unique_ids.add(key.id())
5551
self.assertFalse(key.is_partial())
5652

53+
self.assertEqual(len(unique_ids), 10)
54+
5755

5856
class TestDatastoreSave(TestDatastore):
5957

0 commit comments

Comments
 (0)