Repository: googleapis/google-cloud-python

Issue:
Transaction is left open when error happens in Datastore API
This one is hard to reproduce, but I noticed that when using a transaction with a **with** statement and an error happens inside the transaction itself (a 500 error for example), it appears that the transaction is left open.

``` python
with self.client.transaction():
  # do stuff

# create entities, put, put_multi, etc
```

Exception:

``` python
Traceback (most recent call last):
  File "my_client.py", line 190, in do_stuff
    with self.client.transaction():
  File "site-packages/gcloud/datastore/batch.py", line 275, in __enter__
    self.begin()
  File "site-packages/gcloud/datastore/transaction.py", line 131, in begin
    self._id = self.connection.begin_transaction(self.project)
  File "site-packages/gcloud/datastore/connection.py", line 307, in begin_transaction
    _datastore_pb2.BeginTransactionResponse)
  File "site-packages/gcloud/datastore/connection.py", line 124, in _rpc
    data=request_pb.SerializeToString())
  File "site-packages/gcloud/datastore/connection.py", line 98, in _request
    raise make_exception(headers, error_status.message, use_json=False)
InternalServerError: 500 internal error.
```

After the exception, any operations to the Datastore using the same client have no effect, and no error is reported.
So my guess is that for some reason the transaction was left open and the changes are never committed.



Discussion:
Thanks for the report!

I was able to reproduce this, digging in now.

After just faking the 500

``` diff
diff --git a/google/cloud/datastore/transaction.py b/google/cloud/datastore/transaction.py
index 66ef8bb..0318643 100644
--- a/google/cloud/datastore/transaction.py
+++ b/google/cloud/datastore/transaction.py
@@ -128,6 +128,10 @@ class Transaction(Batch):
         :raises: :class:`ValueError` if the transaction has already begun.
         """
         super(Transaction, self).begin()
+        import httplib2
+        from google.cloud.exceptions import make_exception
+        resp = httplib2.Response({'status': 500})
+        raise make_exception(resp, {'error': {'message': 'internal error'}})
         self._id = self.connection.begin_transaction(self.project)

     def rollback(self):
```

I ran the following:

``` python
>>> from google.cloud import datastore
>>> client = datastore.Client()
>>> q = client.query(kind='Foo')
>>> client.connection._datastore_api
<google.cloud.datastore.connection._DatastoreAPIOverGRPC object at 0x7f4997d6c390>
>>> list(q.fetch())
[]
>>> k = client.key('Foo', 1234)
>>> e = datastore.Entity(key=k)
>>> e['food'] = 'noms'
>>> client.put(e)
>>> list(q.fetch())
[<Entity[{'kind': u'Foo', 'id': 1234L}] {u'food': 'noms'}>]
>>> with client.transaction():
...     e['goodbye'] = 123
...     client.put(e)
... 
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File ".tox/py27/local/lib/python2.7/site-packages/google/cloud/datastore/batch.py", line 275, in __enter__
    self.begin()
  File ".tox/py27/local/lib/python2.7/site-packages/google/cloud/datastore/transaction.py", line 134, in begin
    raise make_exception(resp, {'error': {'message': 'internal error'}})
google.cloud.exceptions.InternalServerError: 500 internal error
>>> list(q.fetch())
[<Entity[{'kind': u'Foo', 'id': 1234L}] {u'food': 'noms'}>]
>>> k2 = client.key('Foo', 123456)
>>> e2 = datastore.Entity(key=k2)
>>> e2['other'] = 'brother'
>>> client.put(e2)
>>> 
>>> list(q.fetch())
[<Entity[{'kind': u'Foo', 'id': 1234L}] {u'food': 'noms'}>]
>>>
>>> # Just use a FRESH client
>>> client = datastore.Client()
>>> client.put(e2)
>>> list(q.fetch())
[<Entity[{'kind': u'Foo', 'id': 1234L}] {u'food': 'noms'}>, <Entity[{'kind': u'Foo', 'id': 123456L}] {u'other': 'brother'}>]
```

The issue is certainly in `Batch.__enter__`, hope to have it sorted out soon.


Resolving patch:
diff --git a/google/cloud/datastore/batch.py b/google/cloud/datastore/batch.py
--- a/google/cloud/datastore/batch.py
+++ b/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):
diff --git a/google/cloud/datastore/client.py b/google/cloud/datastore/client.py
--- a/google/cloud/datastore/client.py
+++ b/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)
diff --git a/google/cloud/datastore/transaction.py b/google/cloud/datastore/transaction.py
--- a/google/cloud/datastore/transaction.py
+++ b/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.


Test patch:
diff --git a/unit_tests/datastore/test_batch.py b/unit_tests/datastore/test_batch.py
--- a/unit_tests/datastore/test_batch.py
+++ b/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,10 +209,19 @@ 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()
@@ -191,12 +229,23 @@ def test_commit(self):
         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):
 
diff --git a/unit_tests/datastore/test_client.py b/unit_tests/datastore/test_client.py
--- a/unit_tests/datastore/test_client.py
+++ b/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)
diff --git a/unit_tests/datastore/test_transaction.py b/unit_tests/datastore/test_transaction.py
--- a/unit_tests/datastore/test_transaction.py
+++ b/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