From e062c8a150026e74bf7fb1500535027f209f370f Mon Sep 17 00:00:00 2001 From: ferponse Date: Tue, 17 Mar 2026 12:47:47 +0100 Subject: [PATCH 1/2] feat(sessions): make prepare_tables() public for eager table initialization MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DatabaseSessionService currently initializes database tables lazily on the first database operation via a private _prepare_tables() method. This causes a noticeable latency spike on the first user request. Rename _prepare_tables() to prepare_tables() so that applications can call it at startup to pay the table-creation cost upfront. The lazy behavior is preserved — each public method still calls prepare_tables() internally. --- .../adk/sessions/database_session_service.py | 17 ++++--- .../sessions/test_session_service.py | 47 +++++++++++++++---- 2 files changed, 48 insertions(+), 16 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index eed1d9eae6..407d73fabc 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -252,12 +252,17 @@ async def _with_session_lock( else: self._session_lock_ref_count[lock_key] = remaining - async def _prepare_tables(self): + async def prepare_tables(self): """Ensure database tables are ready for use. This method is called lazily before each database operation. It checks the DB schema version to use and creates the tables (including setting the schema version metadata) if needed. + + It can also be called eagerly right after construction to pay the + table-creation cost upfront (e.g. during application startup) instead of + on the first database operation. It is safe to call more than once and + is recommended for latency-sensitive applications. """ # Early return if tables are already created if self._tables_created: @@ -323,7 +328,7 @@ async def create_session( # 3. Add the object to the table # 4. Build the session object with generated id # 5. Return the session - await self._prepare_tables() + await self.prepare_tables() schema = self._get_schema_classes() async with self._rollback_on_exception_session() as sql_session: if session_id and await sql_session.get( @@ -398,7 +403,7 @@ async def get_session( session_id: str, config: Optional[GetSessionConfig] = None, ) -> Optional[Session]: - await self._prepare_tables() + await self.prepare_tables() # 1. Get the storage session entry from session table # 2. Get all the events based on session id and filtering config # 3. Convert and return the session @@ -456,7 +461,7 @@ async def get_session( async def list_sessions( self, *, app_name: str, user_id: Optional[str] = None ) -> ListSessionsResponse: - await self._prepare_tables() + await self.prepare_tables() schema = self._get_schema_classes() async with self._rollback_on_exception_session() as sql_session: stmt = select(schema.StorageSession).filter( @@ -506,7 +511,7 @@ async def list_sessions( async def delete_session( self, app_name: str, user_id: str, session_id: str ) -> None: - await self._prepare_tables() + await self.prepare_tables() schema = self._get_schema_classes() async with self._rollback_on_exception_session() as sql_session: stmt = delete(schema.StorageSession).where( @@ -519,7 +524,7 @@ async def delete_session( @override async def append_event(self, session: Session, event: Event) -> Event: - await self._prepare_tables() + await self.prepare_tables() if event.partial: return event diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 9686596b93..854a5ac952 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -1080,10 +1080,10 @@ def _spy_factory(): @pytest.mark.asyncio async def test_concurrent_prepare_tables_no_race_condition(): - """Verifies that concurrent calls to _prepare_tables wait for table creation. + """Verifies that concurrent calls to prepare_tables wait for table creation. Reproduces the race condition from https://github.com/google/adk-python/issues/4445: when concurrent requests - arrive at startup, _prepare_tables must not return before tables exist. + arrive at startup, prepare_tables must not return before tables exist. Previously, the early-return guard checked _db_schema_version (set during schema detection) instead of _tables_created, so a second request could slip through after schema detection but before table creation finished. @@ -1096,7 +1096,7 @@ async def test_concurrent_prepare_tables_no_race_condition(): # Launch several concurrent create_session calls, each with a unique # app_name to avoid IntegrityError on the shared app_states row. - # Each will call _prepare_tables internally. If the race condition + # Each will call prepare_tables internally. If the race condition # exists, some of these will fail because the "sessions" table doesn't # exist yet. num_concurrent = 5 @@ -1114,7 +1114,7 @@ async def test_concurrent_prepare_tables_no_race_condition(): for i, result in enumerate(results): assert not isinstance(result, BaseException), ( f'Concurrent create_session #{i} raised {result!r}; tables were' - ' likely not ready due to the _prepare_tables race condition.' + ' likely not ready due to the prepare_tables race condition.' ) # All sessions should be retrievable. @@ -1133,7 +1133,7 @@ async def test_concurrent_prepare_tables_no_race_condition(): async def test_prepare_tables_serializes_schema_detection_and_creation(): """Verifies schema detection and table creation happen atomically under one lock, so concurrent callers cannot observe a partially-initialized state. - After _prepare_tables completes, both _db_schema_version and _tables_created + After prepare_tables completes, both _db_schema_version and _tables_created must be set. """ service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') @@ -1141,9 +1141,9 @@ async def test_prepare_tables_serializes_schema_detection_and_creation(): assert not service._tables_created assert service._db_schema_version is None - await service._prepare_tables() + await service.prepare_tables() - # Both must be set after a single _prepare_tables call. + # Both must be set after a single prepare_tables call. assert service._tables_created assert service._db_schema_version is not None @@ -1159,17 +1159,17 @@ async def test_prepare_tables_serializes_schema_detection_and_creation(): @pytest.mark.asyncio async def test_prepare_tables_idempotent_after_creation(): - """Calling _prepare_tables multiple times is safe and idempotent. + """Calling prepare_tables multiple times is safe and idempotent. After tables are created, subsequent calls should return immediately via the fast path without errors. """ service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') try: - await service._prepare_tables() + await service.prepare_tables() assert service._tables_created # Call again — should be a no-op via the fast path. - await service._prepare_tables() + await service.prepare_tables() assert service._tables_created # Service should still work. @@ -1181,6 +1181,33 @@ async def test_prepare_tables_idempotent_after_creation(): await service.close() +@pytest.mark.asyncio +async def test_public_prepare_tables_eager_initialization(): + """Calling the public prepare_tables() eagerly initializes tables so that + the first real database operation does not pay the setup cost. + """ + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + try: + # Before calling prepare_tables, tables are not created. + assert not service._tables_created + assert service._db_schema_version is None + + # Eagerly prepare tables via the public API. + await service.prepare_tables() + + # Tables should now be ready. + assert service._tables_created + assert service._db_schema_version is not None + + # Subsequent operations should work without any additional setup cost. + session = await service.create_session( + app_name='app', user_id='user', session_id='s1' + ) + assert session.id == 's1' + finally: + await service.close() + + @pytest.mark.asyncio @pytest.mark.parametrize( 'state_delta, expect_app_lock, expect_user_lock', From dc3b945be542dd76bb5585305dc0f3c545648628 Mon Sep 17 00:00:00 2001 From: ferponse Date: Tue, 17 Mar 2026 12:58:18 +0100 Subject: [PATCH 2/2] refactor(tests): use async context manager in eager initialization test Address review feedback: use `async with` instead of try/finally for DatabaseSessionService resource management, consistent with other tests. --- tests/unittests/sessions/test_session_service.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 854a5ac952..3fa98cdff4 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -1186,8 +1186,7 @@ async def test_public_prepare_tables_eager_initialization(): """Calling the public prepare_tables() eagerly initializes tables so that the first real database operation does not pay the setup cost. """ - service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') - try: + async with DatabaseSessionService('sqlite+aiosqlite:///:memory:') as service: # Before calling prepare_tables, tables are not created. assert not service._tables_created assert service._db_schema_version is None @@ -1204,8 +1203,6 @@ async def test_public_prepare_tables_eager_initialization(): app_name='app', user_id='user', session_id='s1' ) assert session.id == 's1' - finally: - await service.close() @pytest.mark.asyncio