From 6078fc1351497f7e86ad1d870e8290253944b6f1 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 8 Jan 2015 13:20:57 -0500 Subject: [PATCH 1/5] Add 'insert_auto_ids' support to 'Batch'. Follow-on to #509. --- gcloud/datastore/batch.py | 10 +++++- gcloud/datastore/test_batch.py | 59 +++++++++++++++++++++++++++++++--- 2 files changed, 64 insertions(+), 5 deletions(-) diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index 3a9ed176ca78..db985656d3b5 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -72,6 +72,7 @@ def __init__(self, dataset_id=None, connection=None): 'a dataset ID set.') self._mutation = datastore_pb.Mutation() + self._auto_id_entities = [] @property def dataset_id(self): @@ -137,6 +138,9 @@ def put(self, entity): self.dataset_id, key_pb, properties, exclude_from_indexes=exclude, mutation=self.mutation) + if entity.key.is_partial: + self._auto_id_entities.append(entity) + def delete(self, key): """Remember a key to be deleted durring ``commit``. @@ -159,7 +163,11 @@ def commit(self): however it can be called explicitly if you don't want to use a context manager. """ - self.connection.commit(self._dataset_id, self.mutation) + response = self.connection.commit(self._dataset_id, self.mutation) + for new_key_pb, entity in zip(response.insert_auto_id_key, + self._auto_id_entities): + new_id = new_key_pb.path_element[-1].id + entity.key = entity.key.completed_key(new_id) def __enter__(self): return self diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py index 4b3414e50eab..68e7e832b41a 100644 --- a/gcloud/datastore/test_batch.py +++ b/gcloud/datastore/test_batch.py @@ -46,6 +46,7 @@ def test_ctor_explicit(self): self.assertEqual(batch.dataset_id, _DATASET) self.assertEqual(batch.connection, connection) self.assertTrue(isinstance(batch.mutation, Mutation)) + self.assertEqual(batch._auto_id_entities, []) def test_ctor_implicit(self): from gcloud._testing import _Monkey @@ -62,6 +63,7 @@ def test_ctor_implicit(self): self.assertEqual(batch.dataset_id, DATASET_ID) self.assertEqual(batch.connection, CONNECTION) self.assertTrue(isinstance(batch.mutation, Mutation)) + self.assertEqual(batch._auto_id_entities, []) def test_put_entity_wo_key(self): _DATASET = 'DATASET' @@ -70,7 +72,23 @@ def test_put_entity_wo_key(self): self.assertRaises(ValueError, batch.put, _Entity()) - def test_put_entity_w_key(self): + def test_put_entity_w_partial_key(self): + _DATASET = 'DATASET' + _PROPERTIES = {'foo': 'bar'} + connection = _Connection() + batch = self._makeOne(dataset_id=_DATASET, connection=connection) + entity = _Entity(_PROPERTIES) + key = entity.key = _Key(_DATASET) + key._partial = True + + batch.put(entity) + + self.assertEqual( + connection._saved, + (_DATASET, key._key, _PROPERTIES, (), batch.mutation)) + self.assertEqual(batch._auto_id_entities, [entity]) + + def test_put_entity_w_completed_key(self): _DATASET = 'DATASET' _PROPERTIES = {'foo': 'bar'} connection = _Connection() @@ -114,6 +132,22 @@ def test_commit(self): self.assertEqual(connection._committed, (_DATASET, batch.mutation)) + def test_commit_w_auto_id_entities(self): + _DATASET = 'DATASET' + _NEW_ID = 1234 + connection = _Connection(_NEW_ID) + batch = self._makeOne(dataset_id=_DATASET, connection=connection) + entity = _Entity({}) + key = entity.key = _Key(_DATASET) + key._partial = True + batch._auto_id_entities.append(entity) + + batch.commit() + + self.assertEqual(connection._committed, (_DATASET, batch.mutation)) + self.assertFalse(key._partial) + self.assertEqual(key._id, _NEW_ID) + def test_as_context_mgr_wo_error(self): _DATASET = 'DATASET' _PROPERTIES = {'foo': 'bar'} @@ -154,7 +188,19 @@ def test_as_context_mgr_w_error(self): class _CommitResult(object): def __init__(self, *new_keys): - self.insert_auto_id_key = new_keys + self.insert_auto_id_key = [_KeyPB(key) for key in new_keys] + + +class _PathElementPB(object): + + def __init__(self, id): + self.id = id + + +class _KeyPB(object): + + def __init__(self, id): + self.path_element = [_PathElementPB(id)] class _Connection(object): @@ -162,8 +208,8 @@ class _Connection(object): _committed = _saved = _deleted = None _save_result = (False, None) - def __init__(self): - self._commit_result = _CommitResult() + def __init__(self, *new_keys): + self._commit_result = _CommitResult(*new_keys) def save_entity(self, dataset_id, key_pb, properties, exclude_from_indexes=(), mutation=None): @@ -201,3 +247,8 @@ def is_partial(self): def to_protobuf(self): return self._key + + def completed_key(self, new_id): + assert self._partial + self._id = new_id + self._partial = False From d5d28c94277ed97cf6c55fcd21dd9e3886d62759 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 8 Jan 2015 13:29:37 -0500 Subject: [PATCH 2/5] Add 'Batch.add_auto_id_entity' method. FBO the proposed global 'put'. --- gcloud/datastore/batch.py | 19 +++++++++++++++++++ gcloud/datastore/test_batch.py | 21 +++++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index db985656d3b5..bd5493f9d33a 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -108,6 +108,25 @@ def mutation(self): """ return self._mutation + def add_auto_id_entity(self, entity): + """Adds an entity to the list of entities to update with IDs. + + When an entity has a partial key, calling ``save()`` adds an + insert_auto_id entry in the mutation. In order to make sure we + update the Entity once the transaction is committed, we need to + keep track of which entities to update (and the order is + important). + + When you call ``save()`` on an entity inside a transaction, if + the entity has a partial key, it adds itself to the list of + entities to be updated once the transaction is committed by + calling this method. + """ + if not entity.key.is_partial: + raise ValueError("Entity has a completed key") + + self._auto_id_entities.append(entity) + def put(self, entity): """Remember an entity's state to be saved during ``commit``. diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py index 68e7e832b41a..815cd6c3fd0a 100644 --- a/gcloud/datastore/test_batch.py +++ b/gcloud/datastore/test_batch.py @@ -65,6 +65,27 @@ def test_ctor_implicit(self): self.assertTrue(isinstance(batch.mutation, Mutation)) self.assertEqual(batch._auto_id_entities, []) + def test_add_auto_id_entity_w_partial_key(self): + _DATASET = 'DATASET' + connection = _Connection() + batch = self._makeOne(dataset_id=_DATASET, connection=connection) + entity = _Entity() + key = entity.key = _Key(_Entity) + key._partial = True + + batch.add_auto_id_entity(entity) + + self.assertEqual(batch._auto_id_entities, [entity]) + + def test_add_auto_id_entity_w_completed_key(self): + _DATASET = 'DATASET' + connection = _Connection() + batch = self._makeOne(dataset_id=_DATASET, connection=connection) + entity = _Entity() + key = entity.key = _Key(_Entity) + + self.assertRaises(ValueError, batch.add_auto_id_entity, entity) + def test_put_entity_wo_key(self): _DATASET = 'DATASET' connection = _Connection() From de78a69608a00d0351e78f000deca9e680a185f9 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 8 Jan 2015 13:33:03 -0500 Subject: [PATCH 3/5] Chixen bonez. --- gcloud/datastore/test_batch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gcloud/datastore/test_batch.py b/gcloud/datastore/test_batch.py index 815cd6c3fd0a..7768a9df5439 100644 --- a/gcloud/datastore/test_batch.py +++ b/gcloud/datastore/test_batch.py @@ -82,7 +82,7 @@ def test_add_auto_id_entity_w_completed_key(self): connection = _Connection() batch = self._makeOne(dataset_id=_DATASET, connection=connection) entity = _Entity() - key = entity.key = _Key(_Entity) + entity.key = _Key(_Entity) self.assertRaises(ValueError, batch.add_auto_id_entity, entity) From 6f12f16ad2ee031f3e70dd81160e617392735044 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 8 Jan 2015 13:36:13 -0500 Subject: [PATCH 4/5] Docstring fixup. Addresses: https://github.com/GoogleCloudPlatform/gcloud-python/pull/515#commitcomment-9206859 --- gcloud/datastore/batch.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index bd5493f9d33a..ebbf601555f3 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -121,6 +121,11 @@ def add_auto_id_entity(self, entity): the entity has a partial key, it adds itself to the list of entities to be updated once the transaction is committed by calling this method. + + :type entity: :class:`gcloud.datastore.entity.Entity` + :param entity: The entity to be updated with a completed key. + + :raises: ValueError if the entity's key is alread completed. """ if not entity.key.is_partial: raise ValueError("Entity has a completed key") From 5c10b105de35235860a0db5af0c59c6a8ec95256 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Thu, 8 Jan 2015 13:43:14 -0500 Subject: [PATCH 5/5] Add comment re semantics of 'MutationResult.insert_auto_id_keys'. Addresses: https://github.com/GoogleCloudPlatform/gcloud-python/pull/515#discussion_r22670476 --- gcloud/datastore/batch.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/gcloud/datastore/batch.py b/gcloud/datastore/batch.py index ebbf601555f3..e851454cd548 100644 --- a/gcloud/datastore/batch.py +++ b/gcloud/datastore/batch.py @@ -188,6 +188,10 @@ def commit(self): context manager. """ response = self.connection.commit(self._dataset_id, self.mutation) + # If the back-end returns without error, we are guaranteed that + # the response's 'insert_auto_id_key' will match (length and order) + # the request's 'insert_auto_id` entities, which are derived from + # our '_auto_id_entities' (no partial success). for new_key_pb, entity in zip(response.insert_auto_id_key, self._auto_id_entities): new_id = new_key_pb.path_element[-1].id