Skip to content

Commit eb1eaee

Browse files
author
Ferran Pons Serra
committed
feat(sessions): add public prepare_tables() for eager table initialization
DatabaseSessionService currently initializes database tables lazily on the first database operation. This causes a noticeable latency spike on the first user request, which degrades user experience in production applications. Add a public `prepare_tables()` method that delegates to the existing `_prepare_tables()` so that applications can call it at startup time (e.g. during app initialization or health check) to pay the table- creation cost upfront. The lazy behavior is preserved as-is for backward compatibility.
1 parent 22fc332 commit eb1eaee

2 files changed

Lines changed: 38 additions & 0 deletions

File tree

src/google/adk/sessions/database_session_service.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,17 @@ async def _with_session_lock(
252252
else:
253253
self._session_lock_ref_count[lock_key] = remaining
254254

255+
async def prepare_tables(self):
256+
"""Eagerly prepare database tables.
257+
258+
Call this method right after construction to pay the table-creation cost
259+
upfront (e.g. during application startup) instead of on the first
260+
database operation. It is safe to call more than once and is also called
261+
internally before every database operation, so using it is optional but
262+
recommended for latency-sensitive applications.
263+
"""
264+
await self._prepare_tables()
265+
255266
async def _prepare_tables(self):
256267
"""Ensure database tables are ready for use.
257268

tests/unittests/sessions/test_session_service.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1181,6 +1181,33 @@ async def test_prepare_tables_idempotent_after_creation():
11811181
await service.close()
11821182

11831183

1184+
@pytest.mark.asyncio
1185+
async def test_public_prepare_tables_eager_initialization():
1186+
"""Calling the public prepare_tables() eagerly initializes tables so that
1187+
the first real database operation does not pay the setup cost.
1188+
"""
1189+
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
1190+
try:
1191+
# Before calling prepare_tables, tables are not created.
1192+
assert not service._tables_created
1193+
assert service._db_schema_version is None
1194+
1195+
# Eagerly prepare tables via the public API.
1196+
await service.prepare_tables()
1197+
1198+
# Tables should now be ready.
1199+
assert service._tables_created
1200+
assert service._db_schema_version is not None
1201+
1202+
# Subsequent operations should work without any additional setup cost.
1203+
session = await service.create_session(
1204+
app_name='app', user_id='user', session_id='s1'
1205+
)
1206+
assert session.id == 's1'
1207+
finally:
1208+
await service.close()
1209+
1210+
11841211
@pytest.mark.asyncio
11851212
@pytest.mark.parametrize(
11861213
'state_delta, expect_app_lock, expect_user_lock',

0 commit comments

Comments
 (0)