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