From 63ed771d0b5e64fab6587588cc42e95370103968 Mon Sep 17 00:00:00 2001 From: Nikhil Sood <206686452+founderinaction@users.noreply.github.com> Date: Sat, 7 Nov 2026 00:20:18 +0530 Subject: [PATCH] fix(spanner): reset transaction flag when session checkout fails `Database.run_in_transaction` set `_local.transaction_running` before checking out a session, outside the `try`/`finally` that resets it. If `get_session()` raised (for example `ServiceUnavailable` while Spanner is down), the flag stayed set and every later call on that thread failed with "Spanner does not support nested transactions". Check out the session inside the `try`, and only return it to the sessions manager when one was obtained. The async implementation had the same shape and is fixed the same way; the sync module is regenerated output of it. Fixes #18375 --- .../cloud/spanner_v1/_async/database.py | 6 ++- .../google/cloud/spanner_v1/database.py | 6 ++- .../tests/unit/_async/test_database.py | 53 +++++++++++++++++++ .../tests/unit/test_database.py | 50 +++++++++++++++++ 4 files changed, 111 insertions(+), 4 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py index 0fb7de8ec89e..ebcb15e63fb8 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py +++ b/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): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py index 954302146c79..2891a03e3885 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/database.py +++ b/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. diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_database.py b/packages/google-cloud-spanner/tests/unit/_async/test_database.py index bc1f142e05c4..afef3d3866c8 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_database.py +++ b/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 diff --git a/packages/google-cloud-spanner/tests/unit/test_database.py b/packages/google-cloud-spanner/tests/unit/test_database.py index 4b669ccbc433..0b19bc4d86cf 100644 --- a/packages/google-cloud-spanner/tests/unit/test_database.py +++ b/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