Making datastore batch/transaction more robust to failure.

google/cloud/datastore/batch.py

@@ -183,9 +183,13 @@ def put(self, entity):
  :type entity: :class:`google.cloud.datastore.entity.Entity`
  :param entity: the entity to be saved.
 
- :raises: ValueError if entity has no key assigned, or if the key's
+ :raises: :class:`~exceptions.ValueError` if the batch is not in
+ progress, if entity has no key assigned, or if the key's
  ``project`` does not match ours.
  """
+ if self._status != self._IN_PROGRESS:
+ raise ValueError('Batch must be in progress to put()')
+
  if entity.key is None:
  raise ValueError("Entity must have a key")
 
@@ -206,9 +210,13 @@ def delete(self, key):
  :type key: :class:`google.cloud.datastore.key.Key`
  :param key: the key to be deleted.
 
- :raises: ValueError if key is not complete, or if the key's
+ :raises: :class:`~exceptions.ValueError` if the batch is not in
+ progress, if key is not complete, or if the key's
  ``project`` does not match ours.
  """
+ if self._status != self._IN_PROGRESS:
+ raise ValueError('Batch must be in progress to delete()')
+
  if key.is_partial:
  raise ValueError("Key must be complete")
 
@@ -255,7 +263,13 @@ def commit(self):
  This is called automatically upon exiting a with statement,
  however it can be called explicitly if you don't want to use a
  context manager.
+
+ :raises: :class:`~exceptions.ValueError` if the batch is not
+ in progress.
  """
+ if self._status != self._IN_PROGRESS:
+ raise ValueError('Batch must be in progress to commit()')
+
  try:
  self._commit()
  finally:
@@ -267,12 +281,19 @@ def rollback(self):
  Marks the batch as aborted (can't be used again).
 
  Overridden by :class:`google.cloud.datastore.transaction.Transaction`.
+
+ :raises: :class:`~exceptions.ValueError` if the batch is not
+ in progress.
  """
+ if self._status != self._IN_PROGRESS:
+ raise ValueError('Batch must be in progress to rollback()')
+
  self._status = self._ABORTED
 
  def __enter__(self):
- self._client._push_batch(self)
  self.begin()
+ # NOTE: We make sure begin() succeeds before pushing onto the stack.
+ self._client._push_batch(self)
  return self
 
  def __exit__(self, exc_type, exc_val, exc_tb):

google/cloud/datastore/client.py

@@ -348,6 +348,7 @@ def put_multi(self, entities):
 
  if not in_batch:
  current = self.batch()
+ current.begin()
 
  for entity in entities:
  current.put(entity)
@@ -384,6 +385,7 @@ def delete_multi(self, keys):
 
  if not in_batch:
  current = self.batch()
+ current.begin()
 
  for key in keys:
  current.delete(key)

google/cloud/datastore/transaction.py

@@ -90,6 +90,8 @@ class Transaction(Batch):
  :param client: the client used to connect to datastore.
  """
 
+ _status = None
+
  def __init__(self, client):
  super(Transaction, self).__init__(client)
  self._id = None
@@ -125,10 +127,15 @@ def begin(self):
  statement, however it can be called explicitly if you don't want
  to use a context manager.
 
- :raises: :class:`ValueError` if the transaction has already begun.
+ :raises: :class:`~exceptions.ValueError` if the transaction has
+ already begun.
  """
  super(Transaction, self).begin()
- self._id = self.connection.begin_transaction(self.project)
+ try:
+ self._id = self.connection.begin_transaction(self.project)
+ except:
+ self._status = self._ABORTED
+ raise
 
  def rollback(self):
  """Rolls back the current transaction.

unit_tests/datastore/test_batch.py

@@ -68,8 +68,20 @@ def test_put_entity_wo_key(self):
  client = _Client(_PROJECT, connection)
  batch = self._makeOne(client)
 
+ batch.begin()
  self.assertRaises(ValueError, batch.put, _Entity())
 
+ def test_put_entity_wrong_status(self):
+ _PROJECT = 'PROJECT'
+ connection = _Connection()
+ client = _Client(_PROJECT, connection)
+ batch = self._makeOne(client)
+ entity = _Entity()
+ entity.key = _Key('OTHER')
+
+ self.assertEqual(batch._status, batch._INITIAL)
+ self.assertRaises(ValueError, batch.put, entity)
+
  def test_put_entity_w_key_wrong_project(self):
  _PROJECT = 'PROJECT'
  connection = _Connection()
@@ -78,6 +90,7 @@ def test_put_entity_w_key_wrong_project(self):
  entity = _Entity()
  entity.key = _Key('OTHER')
 
+ batch.begin()
  self.assertRaises(ValueError, batch.put, entity)
 
  def test_put_entity_w_partial_key(self):
@@ -90,6 +103,7 @@ def test_put_entity_w_partial_key(self):
  key = entity.key = _Key(_PROJECT)
  key._id = None
 
+ batch.begin()
  batch.put(entity)
 
  mutated_entity = _mutated_pb(self, batch.mutations, 'insert')
@@ -113,6 +127,7 @@ def test_put_entity_w_completed_key(self):
  entity.exclude_from_indexes = ('baz', 'spam')
  key = entity.key = _Key(_PROJECT)
 
+ batch.begin()
  batch.put(entity)
 
  mutated_entity = _mutated_pb(self, batch.mutations, 'upsert')
@@ -129,6 +144,17 @@ def test_put_entity_w_completed_key(self):
  self.assertTrue(spam_values[2].exclude_from_indexes)
  self.assertFalse('frotz' in prop_dict)
 
+ def test_delete_wrong_status(self):
+ _PROJECT = 'PROJECT'
+ connection = _Connection()
+ client = _Client(_PROJECT, connection)
+ batch = self._makeOne(client)
+ key = _Key(_PROJECT)
+ key._id = None
+
+ self.assertEqual(batch._status, batch._INITIAL)
+ self.assertRaises(ValueError, batch.delete, key)
+
  def test_delete_w_partial_key(self):
  _PROJECT = 'PROJECT'
  connection = _Connection()
@@ -137,6 +163,7 @@ def test_delete_w_partial_key(self):
  key = _Key(_PROJECT)
  key._id = None
 
+ batch.begin()
  self.assertRaises(ValueError, batch.delete, key)
 
  def test_delete_w_key_wrong_project(self):
@@ -146,6 +173,7 @@ def test_delete_w_key_wrong_project(self):
  batch = self._makeOne(client)
  key = _Key('OTHER')
 
+ batch.begin()
  self.assertRaises(ValueError, batch.delete, key)
 
  def test_delete_w_completed_key(self):
@@ -155,6 +183,7 @@ def test_delete_w_completed_key(self):
  batch = self._makeOne(client)
  key = _Key(_PROJECT)
 
+ batch.begin()
  batch.delete(key)
 
  mutated_key = _mutated_pb(self, batch.mutations, 'delete')
@@ -180,23 +209,43 @@ def test_rollback(self):
  _PROJECT = 'PROJECT'
  client = _Client(_PROJECT, None)
  batch = self._makeOne(client)
- self.assertEqual(batch._status, batch._INITIAL)
+ batch.begin()
+ self.assertEqual(batch._status, batch._IN_PROGRESS)
  batch.rollback()
  self.assertEqual(batch._status, batch._ABORTED)
 
+ def test_rollback_wrong_status(self):
+ _PROJECT = 'PROJECT'
+ client = _Client(_PROJECT, None)
+ batch = self._makeOne(client)
+
+ self.assertEqual(batch._status, batch._INITIAL)
+ self.assertRaises(ValueError, batch.rollback)
+
  def test_commit(self):
  _PROJECT = 'PROJECT'
  connection = _Connection()
  client = _Client(_PROJECT, connection)
  batch = self._makeOne(client)
 
  self.assertEqual(batch._status, batch._INITIAL)
+ batch.begin()
+ self.assertEqual(batch._status, batch._IN_PROGRESS)
  batch.commit()
  self.assertEqual(batch._status, batch._FINISHED)
 
  self.assertEqual(connection._committed,
  [(_PROJECT, batch._commit_request, None)])
 
+ def test_commit_wrong_status(self):
+ _PROJECT = 'PROJECT'
+ connection = _Connection()
+ client = _Client(_PROJECT, connection)
+ batch = self._makeOne(client)
+
+ self.assertEqual(batch._status, batch._INITIAL)
+ self.assertRaises(ValueError, batch.commit)
+
  def test_commit_w_partial_key_entities(self):
  _PROJECT = 'PROJECT'
  _NEW_ID = 1234
@@ -209,6 +258,8 @@ def test_commit_w_partial_key_entities(self):
  batch._partial_key_entities.append(entity)
 
  self.assertEqual(batch._status, batch._INITIAL)
+ batch.begin()
+ self.assertEqual(batch._status, batch._IN_PROGRESS)
  batch.commit()
  self.assertEqual(batch._status, batch._FINISHED)
 
@@ -295,6 +346,26 @@ def test_as_context_mgr_w_error(self):
  self.assertEqual(mutated_entity.key, key._key)
  self.assertEqual(connection._committed, [])
 
+ def test_as_context_mgr_enter_fails(self):
+ klass = self._getTargetClass()
+
+ class FailedBegin(klass):
+
+ def begin(self):
+ raise RuntimeError
+
+ client = _Client(None, None)
+ self.assertEqual(client._batches, [])
+
+ batch = FailedBegin(client)
+ with self.assertRaises(RuntimeError):
+ # The context manager will never be entered because
+ # of the failure.
+ with batch: # pragma: NO COVER
+ pass
+ # Make sure no batch was added.
+ self.assertEqual(client._batches, [])
+
 
 class _PathElementPB(object):
 

unit_tests/datastore/test_client.py

@@ -960,6 +960,7 @@ def __init__(self, client):
  from google.cloud.datastore.batch import Batch
  self._client = client
  self._batch = Batch(client)
+ self._batch.begin()
 
  def __enter__(self):
  self._client._push_batch(self._batch)
@@ -972,10 +973,12 @@ def __exit__(self, *args):
 class _NoCommitTransaction(object):
 
  def __init__(self, client, transaction_id='TRANSACTION'):
+ from google.cloud.datastore.batch import Batch
  from google.cloud.datastore.transaction import Transaction
  self._client = client
  xact = self._transaction = Transaction(client)
  xact._id = transaction_id
+ Batch.begin(xact)
 
  def __enter__(self):
  self._client._push_batch(self._transaction)

unit_tests/datastore/test_transaction.py

@@ -88,6 +88,19 @@ def test_begin_tombstoned(self):
 
  self.assertRaises(ValueError, xact.begin)
 
+ def test_begin_w_begin_transaction_failure(self):
+ _PROJECT = 'PROJECT'
+ connection = _Connection(234)
+ client = _Client(_PROJECT, connection)
+ xact = self._makeOne(client)
+
+ connection._side_effect = RuntimeError
+ with self.assertRaises(RuntimeError):
+ xact.begin()
+
+ self.assertIsNone(xact.id)
+ self.assertEqual(connection._begun, _PROJECT)
+
  def test_rollback(self):
  _PROJECT = 'PROJECT'
  connection = _Connection(234)
@@ -118,10 +131,10 @@ def test_commit_w_partial_keys(self):
  connection._completed_keys = [_make_key(_KIND, _ID, _PROJECT)]
  client = _Client(_PROJECT, connection)
  xact = self._makeOne(client)
+ xact.begin()
  entity = _Entity()
  xact.put(entity)
  xact._commit_request = commit_request = object()
- xact.begin()
  xact.commit()
  self.assertEqual(connection._committed,
  (_PROJECT, commit_request, 234))
@@ -176,7 +189,10 @@ def _make_key(kind, id_, project):
 
 class _Connection(object):
  _marker = object()
- _begun = _rolled_back = _committed = None
+ _begun = None
+ _rolled_back = None
+ _committed = None
+ _side_effect = None
 
  def __init__(self, xact_id=123):
  self._xact_id = xact_id
@@ -185,7 +201,10 @@ def __init__(self, xact_id=123):
 
  def begin_transaction(self, project):
  self._begun = project
- return self._xact_id
+ if self._side_effect is None:
+ return self._xact_id
+ else:
+ raise self._side_effect
 
  def rollback(self, project, transaction_id):
  self._rolled_back = project, transaction_id