From 63a8eba53f2cb07625eb7cd111ff767e8e0030fa Mon Sep 17 00:00:00 2001 From: Liang Wu Date: Fri, 30 Jan 2026 11:09:23 -0800 Subject: [PATCH] fix: Ensure database sessions are always rolled back on errors Fixes issue#3328 Co-authored-by: Liang Wu PiperOrigin-RevId: 863314939 --- .../adk/sessions/database_session_service.py | 31 ++- .../sessions/test_session_service.py | 213 ++++++++++++++++++ 2 files changed, 238 insertions(+), 6 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index afc09b3e..92b37229 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -14,11 +14,13 @@ from __future__ import annotations import asyncio +from contextlib import asynccontextmanager import copy from datetime import datetime from datetime import timezone import logging from typing import Any +from typing import AsyncIterator from typing import Optional from sqlalchemy import delete @@ -156,6 +158,23 @@ class DatabaseSessionService(BaseSessionService): def _get_schema_classes(self) -> _SchemaClasses: return _SchemaClasses(self._db_schema_version) + @asynccontextmanager + async def _rollback_on_exception_session( + self, + ) -> AsyncIterator[DatabaseSessionFactory]: + """Yields a database session with guaranteed rollback on errors. + + On normal exit the caller is responsible for committing; on any exception + the transaction is explicitly rolled back before the error propagates, + preventing connection-pool exhaustion from lingering invalid transactions. + """ + async with self.database_session_factory() as sql_session: + try: + yield sql_session + except BaseException: + await sql_session.rollback() + raise + async def _prepare_tables(self): """Ensure database tables are ready for use. @@ -204,7 +223,7 @@ class DatabaseSessionService(BaseSessionService): self._tables_created = True if self._db_schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION: - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: # Check if schema version is set, if not, set it to the latest # version stmt = select(StorageMetadata).where( @@ -236,7 +255,7 @@ class DatabaseSessionService(BaseSessionService): # 5. Return the session await self._prepare_tables() schema = self._get_schema_classes() - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: if session_id and await sql_session.get( schema.StorageSession, (app_name, user_id, session_id) ): @@ -313,7 +332,7 @@ class DatabaseSessionService(BaseSessionService): # 2. Get all the events based on session id and filtering config # 3. Convert and return the session schema = self._get_schema_classes() - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: storage_session = await sql_session.get( schema.StorageSession, (app_name, user_id, session_id) ) @@ -368,7 +387,7 @@ class DatabaseSessionService(BaseSessionService): ) -> ListSessionsResponse: await self._prepare_tables() schema = self._get_schema_classes() - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: stmt = select(schema.StorageSession).filter( schema.StorageSession.app_name == app_name ) @@ -418,7 +437,7 @@ class DatabaseSessionService(BaseSessionService): ) -> None: await self._prepare_tables() schema = self._get_schema_classes() - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: stmt = delete(schema.StorageSession).where( schema.StorageSession.app_name == app_name, schema.StorageSession.user_id == user_id, @@ -440,7 +459,7 @@ class DatabaseSessionService(BaseSessionService): # 2. Update session attributes based on event config # 3. Store event to table schema = self._get_schema_classes() - async with self.database_session_factory() as sql_session: + async with self._rollback_on_exception_session() as sql_session: storage_session = await sql_session.get( schema.StorageSession, (session.app_name, session.user_id, session.id) ) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 75645922..a305f53c 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -643,3 +643,216 @@ async def test_partial_events_are_not_persisted(session_service): app_name=app_name, user_id=user_id, session_id=session.id ) assert len(session_got.events) == 0 + + +# --------------------------------------------------------------------------- +# Rollback tests – verify _rollback_on_exception_session explicitly rolls back +# on errors +# --------------------------------------------------------------------------- +class _RollbackSpySession: + """Wraps an AsyncSession to spy on rollback() and optionally fail commit().""" + + def __init__(self, real_session, *, fail_commit=False): + self._real = real_session + self._fail_commit = fail_commit + self.rollback_called = False + + async def __aenter__(self): + self._real = await self._real.__aenter__() + return self + + async def __aexit__(self, *args): + return await self._real.__aexit__(*args) + + async def commit(self): + if self._fail_commit: + raise RuntimeError('simulated commit failure') + return await self._real.commit() + + async def rollback(self): + self.rollback_called = True + return await self._real.rollback() + + def __getattr__(self, name): + return getattr(self._real, name) + + +@pytest.mark.asyncio +async def test_create_session_calls_rollback_on_commit_failure(): + """Verifies that a commit failure during create_session triggers an explicit + rollback() call via _rollback_on_exception_session, not just a close().""" + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + try: + # Ensure tables are initialized. + await service.create_session( + app_name='app', user_id='user', session_id='good' + ) + + original_factory = service.database_session_factory + spy_sessions = [] + + def _spy_factory(): + spy = _RollbackSpySession(original_factory(), fail_commit=True) + spy_sessions.append(spy) + return spy + + service.database_session_factory = _spy_factory + + with pytest.raises(RuntimeError, match='simulated commit failure'): + await service.create_session( + app_name='app', user_id='user', session_id='should_fail' + ) + + # The key assertion: rollback() must have been called explicitly. + assert len(spy_sessions) == 1 + assert spy_sessions[0].rollback_called, ( + 'rollback() was not called – _rollback_on_exception_session is not' + ' protecting this path' + ) + + # Restore and verify the failed session was not persisted. + service.database_session_factory = original_factory + assert ( + await service.get_session( + app_name='app', user_id='user', session_id='should_fail' + ) + is None + ) + finally: + await service.close() + + +@pytest.mark.asyncio +async def test_append_event_calls_rollback_on_commit_failure(): + """Verifies that a commit failure during append_event triggers an explicit + rollback() call via _rollback_on_exception_session.""" + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + try: + session = await service.create_session( + app_name='app', user_id='user', session_id='s1' + ) + + # Successfully append one event first. + event1 = Event( + invocation_id='inv1', + author='user', + actions=EventActions(state_delta={'key1': 'value1'}), + ) + await service.append_event(session, event1) + + original_factory = service.database_session_factory + spy_sessions = [] + + def _spy_factory(): + spy = _RollbackSpySession(original_factory(), fail_commit=True) + spy_sessions.append(spy) + return spy + + service.database_session_factory = _spy_factory + + event2 = Event( + invocation_id='inv2', + author='user', + actions=EventActions(state_delta={'key2': 'value2'}), + ) + with pytest.raises(RuntimeError, match='simulated commit failure'): + await service.append_event(session, event2) + + assert len(spy_sessions) == 1 + assert spy_sessions[0].rollback_called, ( + 'rollback() was not called – _rollback_on_exception_session is not' + ' protecting this path' + ) + + # Restore and verify only the first event was persisted. + service.database_session_factory = original_factory + got = await service.get_session( + app_name='app', user_id='user', session_id='s1' + ) + assert len(got.events) == 1 + assert got.events[0].invocation_id == 'inv1' + finally: + await service.close() + + +@pytest.mark.asyncio +async def test_delete_session_calls_rollback_on_commit_failure(): + """Verifies that a commit failure during delete_session triggers an explicit + rollback() call via _rollback_on_exception_session.""" + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + try: + await service.create_session( + app_name='app', user_id='user', session_id='s1' + ) + + original_factory = service.database_session_factory + spy_sessions = [] + + def _spy_factory(): + spy = _RollbackSpySession(original_factory(), fail_commit=True) + spy_sessions.append(spy) + return spy + + service.database_session_factory = _spy_factory + + with pytest.raises(RuntimeError, match='simulated commit failure'): + await service.delete_session( + app_name='app', user_id='user', session_id='s1' + ) + + assert len(spy_sessions) == 1 + assert spy_sessions[0].rollback_called, ( + 'rollback() was not called – _rollback_on_exception_session is not' + ' protecting this path' + ) + + # Restore and verify the session still exists (delete was rolled back). + service.database_session_factory = original_factory + got = await service.get_session( + app_name='app', user_id='user', session_id='s1' + ) + assert got is not None + finally: + await service.close() + + +@pytest.mark.asyncio +async def test_service_recovers_after_multiple_failures(): + """After several consecutive commit failures, every single one must trigger + a rollback() call and the service must remain functional afterward.""" + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + try: + await service.create_session( + app_name='app', user_id='user', session_id='seed' + ) + + original_factory = service.database_session_factory + spy_sessions = [] + + def _spy_factory(): + spy = _RollbackSpySession(original_factory(), fail_commit=True) + spy_sessions.append(spy) + return spy + + service.database_session_factory = _spy_factory + + num_failures = 5 + for i in range(num_failures): + with pytest.raises(RuntimeError, match='simulated commit failure'): + await service.create_session( + app_name='app', user_id='user', session_id=f'fail_{i}' + ) + + # Every failure must have triggered a rollback. + assert len(spy_sessions) == num_failures + for i, spy in enumerate(spy_sessions): + assert spy.rollback_called, f'rollback() was not called on failure #{i}' + + # Restore and verify the service is still healthy. + service.database_session_factory = original_factory + session = await service.create_session( + app_name='app', user_id='user', session_id='recovered' + ) + assert session.id == 'recovered' + finally: + await service.close()