diff --git a/gcloud/datastore/dataset.py b/gcloud/datastore/dataset.py index c527c34aab1fa..6760a742e5a7c 100644 --- a/gcloud/datastore/dataset.py +++ b/gcloud/datastore/dataset.py @@ -182,3 +182,29 @@ def get_entities(self, keys, missing=None, deferred=None): entities.append(helpers.entity_from_protobuf( entity_pb, dataset=self)) return entities + + def allocate_ids(self, incomplete_key, num_ids): + """Allocates a list of IDs from a partial key. + + :type incomplete_key: A :class:`gcloud.datastore.key.Key` + :param incomplete_key: Partial key to use as base for allocated IDs. + + :type num_ids: A :class:`int`. + :param num_ids: The number of IDs to allocate. + + :rtype: list of :class:`gcloud.datastore.key.Key` + :return: The (complete) keys allocated with `incomplete_key` as root. + :raises: `ValueError` if `incomplete_key` is not a partial key. + """ + if not incomplete_key.is_partial(): + raise ValueError(('Key is not partial.', incomplete_key)) + + incomplete_key_pb = incomplete_key.to_protobuf() + incomplete_key_pbs = [incomplete_key_pb] * num_ids + + allocated_key_pbs = self.connection().allocate_ids( + self.id(), incomplete_key_pbs) + allocated_ids = [allocated_key_pb.path_element[-1].id + for allocated_key_pb in allocated_key_pbs] + return [incomplete_key.id(allocated_id) + for allocated_id in allocated_ids] diff --git a/gcloud/datastore/test_dataset.py b/gcloud/datastore/test_dataset.py index a59e4b307b178..f155198f963ee 100644 --- a/gcloud/datastore/test_dataset.py +++ b/gcloud/datastore/test_dataset.py @@ -213,6 +213,40 @@ def test_get_entities_hit(self): self.assertEqual(list(result), ['foo']) self.assertEqual(result['foo'], 'Foo') + def test_allocate_ids(self): + from gcloud.datastore.test_entity import _Key + + INCOMPLETE_KEY = _Key() + PROTO_ID = object() + INCOMPLETE_KEY._key = _KeyProto(PROTO_ID) + INCOMPLETE_KEY._partial = True + + CONNECTION = _Connection() + NUM_IDS = 2 + DATASET_ID = 'foo' + DATASET = self._makeOne(DATASET_ID, connection=CONNECTION) + result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS) + + # Check the IDs returned match _PathElementProto. + self.assertEqual(result, range(NUM_IDS)) + + # Check connection is called correctly. + self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID) + self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS) + + # Check the IDs passed to Connection.allocate_ids. + key_paths = [key_pb.path_element[-1].id + for key_pb in CONNECTION._called_key_pbs] + self.assertEqual(key_paths, [PROTO_ID] * NUM_IDS) + + def test_allocate_ids_with_complete(self): + from gcloud.datastore.test_entity import _Key + + COMPLETE_KEY = _Key() + DATASET = self._makeOne(None) + self.assertRaises(ValueError, DATASET.allocate_ids, + COMPLETE_KEY, 2) + class _Connection(object): _called_with = None @@ -230,3 +264,21 @@ def lookup(self, **kw): if deferred is not None: deferred.extend(self._deferred) return self._result + + def allocate_ids(self, dataset_id, key_pbs): + self._called_dataset_id = dataset_id + self._called_key_pbs = key_pbs + num_pbs = len(key_pbs) + return [_KeyProto(i) for i in xrange(num_pbs)] + + +class _PathElementProto(object): + + def __init__(self, _id): + self.id = _id + + +class _KeyProto(object): + + def __init__(self, id_): + self.path_element = [_PathElementProto(id_)] diff --git a/gcloud/datastore/test_entity.py b/gcloud/datastore/test_entity.py index 164a69aeb003d..b2cafe26485d3 100644 --- a/gcloud/datastore/test_entity.py +++ b/gcloud/datastore/test_entity.py @@ -240,6 +240,10 @@ class _Key(object): _partial = False _path = None + def id(self, id_to_set): + self._called_id = id_to_set + return id_to_set + def to_protobuf(self): return self._key diff --git a/regression/datastore.py b/regression/datastore.py index 811e5ac6b24ee..c8670ae32892e 100644 --- a/regression/datastore.py +++ b/regression/datastore.py @@ -42,18 +42,16 @@ class TestDatastoreAllocateIDs(TestDatastore): def test_allocate_ids(self): incomplete_key = datastore.key.Key(path=[{'kind': 'Kind'}]) - incomplete_key_pb = incomplete_key.to_protobuf() - incomplete_key_pbs = [incomplete_key_pb] * 10 - - connection = self.dataset.connection() - allocated_key_pbs = connection.allocate_ids(self.dataset.id(), - incomplete_key_pbs) - allocated_keys = [datastore.helpers.key_from_protobuf(key_pb) - for key_pb in allocated_key_pbs] + allocated_keys = self.dataset.allocate_ids(incomplete_key, 10) self.assertEqual(len(allocated_keys), 10) + + unique_ids = set() for key in allocated_keys: + unique_ids.add(key.id()) self.assertFalse(key.is_partial()) + self.assertEqual(len(unique_ids), 10) + class TestDatastoreSave(TestDatastore):