Ensure 'session._transaction' is cleared after 'run_in_transaction' succeeds (googleapis#3437)

tseaver · web-flow · commit f6dab412f5b7 · 2017-05-17T17:01:00.000-04:00
* Ensure 'session._transaction' is cleared after 'run_in_transaction' succeeds. Closes googleapis#3434.
diff --git a/spanner/google/cloud/spanner/session.py b/spanner/google/cloud/spanner/session.py
@@ -311,7 +311,9 @@ def run_in_transaction(self, func, *args, **kw):
                 _delay_until_retry(exc, deadline)
                 del self._transaction
             else:
-                return txn.committed
+                committed = txn.committed
+                del self._transaction
+                return committed
 
 
 # pylint: disable=misplaced-bare-raise
diff --git a/spanner/tests/system/test_system.py b/spanner/tests/system/test_system.py
@@ -323,6 +323,23 @@ def _unit_of_work(transaction, test):
         rows = list(self._db.execute_sql(self.SQL))
         self._check_row_data(rows)
 
+    def test_db_run_in_transaction_twice(self):
+        retry = RetryInstanceState(_has_all_ddl)
+        retry(self._db.reload)()
+
+        with self._db.batch() as batch:
+            batch.delete(self.TABLE, self.ALL)
+
+        def _unit_of_work(transaction, test):
+            transaction.insert_or_update(
+                test.TABLE, test.COLUMNS, test.ROW_DATA)
+
+        self._db.run_in_transaction(_unit_of_work, test=self)
+        self._db.run_in_transaction(_unit_of_work, test=self)
+
+        rows = list(self._db.execute_sql(self.SQL))
+        self._check_row_data(rows)
+
 
 class TestSessionAPI(unittest.TestCase, _TestData):
     ALL_TYPES_TABLE = 'all_types'
diff --git a/spanner/tests/unit/test_session.py b/spanner/tests/unit/test_session.py
@@ -503,6 +503,7 @@ def unit_of_work(txn, *args, **kw):
             unit_of_work, 'abc', some_arg='def')
 
         self.assertEqual(committed, now)
+        self.assertIsNone(session._transaction)
         self.assertEqual(len(called_with), 1)
         txn, args, kw = called_with[0]
         self.assertIsInstance(txn, Transaction)
