feat(sessions): add public prepare_tables() for eager initialization

Merge #4858

Fixes #4857

## Summary

- Rename `_prepare_tables()` to `prepare_tables()`, making it part of the public API so applications can eagerly initialize database tables at startup and avoid the latency spike on the first user request.
- All internal callers (`create_session`, `get_session`, `list_sessions`, `delete_session`, `append_event`) updated to call the now-public method. Lazy behavior is fully preserved.
- Existing tests updated to use the public method name.
- New test (`test_public_prepare_tables_eager_initialization`) verifies eager initialization works correctly.

## Motivation

In our production application using `DatabaseSessionService` with PostgreSQL, the first user query after a restart is noticeably slower because `_prepare_tables()` runs schema detection + `metadata.create_all()` inline. By making the method public, applications can call it during startup (e.g. in a FastAPI lifespan handler) to eliminate this cold-start penalty:

```python
session_service = DatabaseSessionService(db_url=db_url)
await session_service.prepare_tables()  # eager initialization at startup
```

## Test plan

- [x] New test `test_public_prepare_tables_eager_initialization` passes
- [x] All existing `prepare_tables` tests still pass (4/4)
- [x] Fully backward compatible — lazy initialization still works for applications that don't call `prepare_tables()` explicitly

Co-authored-by: Liang Wu <wuliang@google.com>
COPYBARA_INTEGRATE_REVIEW=#4858 from ferponse:feat/eager-table-preparation aa911af
PiperOrigin-RevId: 934013793

Author: 2 people

src/google/adk/sessions/database_session_service.py

@@ -330,12 +330,17 @@ async def _with_session_lock(
         else:
           self._session_lock_ref_count[lock_key] = remaining
 
-  async def _prepare_tables(self) -> None:
+  async def prepare_tables(self) -> None:
     """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:
@@ -434,7 +439,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()
     has_user_provided_id = session_id is not None
     if session_id is None:
       session_id = platform_uuid.new_uuid()
@@ -510,7 +515,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
@@ -574,7 +579,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(
         read_only=True
@@ -631,7 +636,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(
@@ -646,7 +651,7 @@ async def delete_session(
   async def get_user_state(
       self, *, app_name: str, user_id: str
   ) -> dict[str, Any]:
-    await self._prepare_tables()
+    await self.prepare_tables()
     schema = self._get_schema_classes()
     async with self._rollback_on_exception_session(
         read_only=True
@@ -660,7 +665,7 @@ async def get_user_state(
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
-    await self._prepare_tables()
+    await self.prepare_tables()
     if event.partial:
       return event
 

tests/unittests/sessions/test_session_service.py

@@ -285,7 +285,7 @@ async def test_get_empty_session(session_service):
 @pytest.mark.asyncio
 async def test_database_session_service_get_session_uses_read_only_factory():
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
-  service._prepare_tables = mock.AsyncMock()
+  service.prepare_tables = mock.AsyncMock()
 
   read_only_session = mock.AsyncMock()
   read_only_session.get = mock.AsyncMock(return_value=None)
@@ -315,7 +315,7 @@ async def fake_read_only_session():
 @pytest.mark.asyncio
 async def test_database_session_service_list_sessions_uses_read_only_factory():
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
-  service._prepare_tables = mock.AsyncMock()
+  service.prepare_tables = mock.AsyncMock()
 
   read_only_session = mock.AsyncMock()
   empty_result = mock.Mock()
@@ -1281,7 +1281,7 @@ def __getattr__(self, name):
 async def test_append_event_reads_storage_revision_before_commit():
   """append_event captures session revision before commit completes."""
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
-  await service._prepare_tables()
+  await service.prepare_tables()
   schema = service._get_schema_classes()
   original_get_update_timestamp = schema.StorageSession.get_update_timestamp
   original_get_update_marker = schema.StorageSession.get_update_marker
@@ -1339,7 +1339,7 @@ def _spy_factory():
 async def test_create_session_reads_storage_revision_before_commit():
   """create_session captures session revision before commit completes."""
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
-  await service._prepare_tables()
+  await service.prepare_tables()
   schema = service._get_schema_classes()
   original_get_update_timestamp = schema.StorageSession.get_update_timestamp
   original_get_update_marker = schema.StorageSession.get_update_marker
@@ -1470,10 +1470,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.
@@ -1486,7 +1486,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
@@ -1504,7 +1504,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.
@@ -1523,17 +1523,17 @@ 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:')
   try:
     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
 
@@ -1552,7 +1552,7 @@ async def test_get_or_create_state_returns_existing_row():
   """_get_or_create_state returns an existing row without inserting."""
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
   try:
-    await service._prepare_tables()
+    await service.prepare_tables()
     schema = service._get_schema_classes()
 
     # Pre-create the app_state row.
@@ -1579,7 +1579,7 @@ async def test_get_or_create_state_creates_new_row():
   """_get_or_create_state creates a row when none exists."""
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
   try:
-    await service._prepare_tables()
+    await service.prepare_tables()
     schema = service._get_schema_classes()
 
     async with service.database_session_factory() as sql_session:
@@ -1612,7 +1612,7 @@ async def test_get_or_create_state_handles_race_condition():
   """
   service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
   try:
-    await service._prepare_tables()
+    await service.prepare_tables()
     schema = service._get_schema_classes()
 
     # Pre-create the row to guarantee the INSERT will fail.
@@ -1680,17 +1680,17 @@ async def test_create_session_sequential_same_app_name():
 
 @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.
@@ -1702,6 +1702,30 @@ 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.
+  """
+  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
+
+    # 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'
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     'state_delta, expect_app_lock, expect_user_lock',