Add 'insert_auto_ids' support to 'Batch'.

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):
@@ -107,6 +108,30 @@ 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.
+
+ :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")
+
+ self._auto_id_entities.append(entity)
+
  def put(self, entity):
  """Remember an entity's state to be saved during ``commit``.
 
@@ -137,6 +162,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 +187,15 @@ 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)
+ # 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
+ entity.key = entity.key.completed_key(new_id)
 
  def __enter__(self):
  return self

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,28 @@ 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_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()
+ entity.key = _Key(_Entity)
+
+ self.assertRaises(ValueError, batch.add_auto_id_entity, entity)
 
  def test_put_entity_wo_key(self):
  _DATASET = 'DATASET'
@@ -70,7 +93,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 +153,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,16 +209,28 @@ 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):
  _marker = 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 +268,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