fix(spanner): reset transaction flag when session checkout fails

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py

@@ -1233,14 +1233,16 @@ async def run_in_transaction(self, func, *args, **kw):
             # Check out a session and run the function in a transaction; once
             # done, flip the sanity check bit back and return the session.
             transaction_type = TransactionType.READ_WRITE
-            session = await self._sessions_manager.get_session(transaction_type)
+            session = None
 
             try:
+                session = await self._sessions_manager.get_session(transaction_type)
                 return await session.run_in_transaction(func, *args, **kw)
 
             finally:
                 self._local.transaction_running = False
-                await self._sessions_manager.put_session(session)
+                if session is not None:
+                    await self._sessions_manager.put_session(session)
 
     @CrossSync.convert
     async def restore(self, source):

packages/google-cloud-spanner/google/cloud/spanner_v1/database.py

@@ -1090,12 +1090,14 @@ def run_in_transaction(self, func, *args, **kw):
                 raise RuntimeError("Spanner does not support nested transactions.")
             self._local.transaction_running = True
             transaction_type = TransactionType.READ_WRITE
-            session = self._sessions_manager.get_session(transaction_type)
+            session = None
             try:
+                session = self._sessions_manager.get_session(transaction_type)
                 return session.run_in_transaction(func, *args, **kw)
             finally:
                 self._local.transaction_running = False
-                self._sessions_manager.put_session(session)
+                if session is not None:
+                    self._sessions_manager.put_session(session)
 
     def restore(self, source):
         """Restore from a backup to this database.

packages/google-cloud-spanner/tests/unit/_async/test_database.py

@@ -2019,6 +2019,59 @@ def _unit_of_work(txn, *args, **kwargs):
 
             self.assertEqual(committed, NOW)
 
+    @CrossSync.pytest
+    async def test_run_in_transaction_session_error(self):
+        from google.api_core.exceptions import ServiceUnavailable
+
+        instance = _Instance(self.INSTANCE_NAME, client=_Client())
+        database = await self._make_one(self.DATABASE_ID, instance, pool=_Pool())
+        session = _Session(run_transaction_function=True)
+        session._committed = object()
+        unit_of_work = mock.AsyncMock(return_value=session._committed)
+
+        with mock.patch.object(database, "_sessions_manager") as sessions_manager:
+            sessions_manager.get_session = mock.AsyncMock(
+                side_effect=[ServiceUnavailable("Session unavailable"), session]
+            )
+            sessions_manager.put_session = mock.AsyncMock()
+
+            with self.assertRaises(ServiceUnavailable):
+                await database.run_in_transaction(unit_of_work)
+
+            unit_of_work.assert_not_called()
+            sessions_manager.put_session.assert_not_called()
+
+            result = await database.run_in_transaction(unit_of_work)
+
+            self.assertIs(result, session._committed)
+            unit_of_work.assert_awaited_once_with(mock.ANY)
+            self.assertFalse(database._local.transaction_running)
+            self.assertEqual(
+                sessions_manager.get_session.await_args_list,
+                [mock.call(TransactionType.READ_WRITE)] * 2,
+            )
+            sessions_manager.put_session.assert_awaited_once_with(session)
+
+    @CrossSync.pytest
+    async def test_run_in_transaction_callback_error(self):
+        instance = _Instance(self.INSTANCE_NAME, client=_Client())
+        database = await self._make_one(self.DATABASE_ID, instance, pool=_Pool())
+        session = _Session(run_transaction_function=True)
+        error = ValueError("Callback failed")
+        unit_of_work = mock.AsyncMock(side_effect=error)
+
+        with mock.patch.object(database, "_sessions_manager") as sessions_manager:
+            sessions_manager.get_session = mock.AsyncMock(return_value=session)
+            sessions_manager.put_session = mock.AsyncMock()
+
+            with self.assertRaises(ValueError) as raised:
+                await database.run_in_transaction(unit_of_work)
+
+            self.assertIs(raised.exception, error)
+            unit_of_work.assert_awaited_once_with(mock.ANY)
+            self.assertFalse(database._local.transaction_running)
+            sessions_manager.put_session.assert_awaited_once_with(session)
+
     @CrossSync.pytest
     async def test_run_in_transaction_tracing(self):
         from opentelemetry.sdk.trace import TracerProvider

packages/google-cloud-spanner/tests/unit/test_database.py

@@ -1905,6 +1905,56 @@ def _unit_of_work(txn, *args, **kwargs):
 
             self.assertEqual(committed, NOW)
 
+    def test_run_in_transaction_session_error(self):
+        from google.api_core.exceptions import ServiceUnavailable
+
+        instance = _Instance(self.INSTANCE_NAME, client=_Client())
+        database = self._make_one(self.DATABASE_ID, instance, pool=_Pool())
+        session = _Session(run_transaction_function=True)
+        session._committed = object()
+        unit_of_work = mock.Mock(return_value=session._committed)
+
+        with mock.patch.object(database, "_sessions_manager") as sessions_manager:
+            sessions_manager.get_session.side_effect = [
+                ServiceUnavailable("Session unavailable"),
+                session,
+            ]
+
+            with self.assertRaises(ServiceUnavailable):
+                database.run_in_transaction(unit_of_work)
+
+            unit_of_work.assert_not_called()
+            sessions_manager.put_session.assert_not_called()
+
+            result = database.run_in_transaction(unit_of_work)
+
+            self.assertIs(result, session._committed)
+            unit_of_work.assert_called_once_with(mock.ANY)
+            self.assertFalse(database._local.transaction_running)
+            self.assertEqual(
+                sessions_manager.get_session.call_args_list,
+                [mock.call(TransactionType.READ_WRITE)] * 2,
+            )
+            sessions_manager.put_session.assert_called_once_with(session)
+
+    def test_run_in_transaction_callback_error(self):
+        instance = _Instance(self.INSTANCE_NAME, client=_Client())
+        database = self._make_one(self.DATABASE_ID, instance, pool=_Pool())
+        session = _Session(run_transaction_function=True)
+        error = ValueError("Callback failed")
+        unit_of_work = mock.Mock(side_effect=error)
+
+        with mock.patch.object(database, "_sessions_manager") as sessions_manager:
+            sessions_manager.get_session.return_value = session
+
+            with self.assertRaises(ValueError) as raised:
+                database.run_in_transaction(unit_of_work)
+
+            self.assertIs(raised.exception, error)
+            unit_of_work.assert_called_once_with(mock.ANY)
+            self.assertFalse(database._local.transaction_running)
+            sessions_manager.put_session.assert_called_once_with(session)
+
     def test_run_in_transaction_tracing(self):
         from opentelemetry.sdk.trace import TracerProvider
         from opentelemetry.sdk.trace.export import SimpleSpanProcessor