Skip to content

Commit e27aad7

Browse files
author
Ferran Pons Serra
committed
feat(sessions): make prepare_tables() public for eager table initialization
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.
1 parent 22fc332 commit e27aad7

2 files changed

Lines changed: 48 additions & 16 deletions

File tree

src/google/adk/sessions/database_session_service.py

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -252,12 +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):
255+
async def prepare_tables(self):
256256
"""Ensure database tables are ready for use.
257257
258258
This method is called lazily before each database operation. It checks the
259259
DB schema version to use and creates the tables (including setting the
260260
schema version metadata) if needed.
261+
262+
It can also be called eagerly right after construction to pay the
263+
table-creation cost upfront (e.g. during application startup) instead of
264+
on the first database operation. It is safe to call more than once and
265+
is recommended for latency-sensitive applications.
261266
"""
262267
# Early return if tables are already created
263268
if self._tables_created:
@@ -323,7 +328,7 @@ async def create_session(
323328
# 3. Add the object to the table
324329
# 4. Build the session object with generated id
325330
# 5. Return the session
326-
await self._prepare_tables()
331+
await self.prepare_tables()
327332
schema = self._get_schema_classes()
328333
async with self._rollback_on_exception_session() as sql_session:
329334
if session_id and await sql_session.get(
@@ -398,7 +403,7 @@ async def get_session(
398403
session_id: str,
399404
config: Optional[GetSessionConfig] = None,
400405
) -> Optional[Session]:
401-
await self._prepare_tables()
406+
await self.prepare_tables()
402407
# 1. Get the storage session entry from session table
403408
# 2. Get all the events based on session id and filtering config
404409
# 3. Convert and return the session
@@ -456,7 +461,7 @@ async def get_session(
456461
async def list_sessions(
457462
self, *, app_name: str, user_id: Optional[str] = None
458463
) -> ListSessionsResponse:
459-
await self._prepare_tables()
464+
await self.prepare_tables()
460465
schema = self._get_schema_classes()
461466
async with self._rollback_on_exception_session() as sql_session:
462467
stmt = select(schema.StorageSession).filter(
@@ -506,7 +511,7 @@ async def list_sessions(
506511
async def delete_session(
507512
self, app_name: str, user_id: str, session_id: str
508513
) -> None:
509-
await self._prepare_tables()
514+
await self.prepare_tables()
510515
schema = self._get_schema_classes()
511516
async with self._rollback_on_exception_session() as sql_session:
512517
stmt = delete(schema.StorageSession).where(
@@ -519,7 +524,7 @@ async def delete_session(
519524

520525
@override
521526
async def append_event(self, session: Session, event: Event) -> Event:
522-
await self._prepare_tables()
527+
await self.prepare_tables()
523528
if event.partial:
524529
return event
525530

tests/unittests/sessions/test_session_service.py

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1080,10 +1080,10 @@ def _spy_factory():
10801080

10811081
@pytest.mark.asyncio
10821082
async def test_concurrent_prepare_tables_no_race_condition():
1083-
"""Verifies that concurrent calls to _prepare_tables wait for table creation.
1083+
"""Verifies that concurrent calls to prepare_tables wait for table creation.
10841084
Reproduces the race condition from
10851085
https://github.com/google/adk-python/issues/4445: when concurrent requests
1086-
arrive at startup, _prepare_tables must not return before tables exist.
1086+
arrive at startup, prepare_tables must not return before tables exist.
10871087
Previously, the early-return guard checked _db_schema_version (set during
10881088
schema detection) instead of _tables_created, so a second request could
10891089
slip through after schema detection but before table creation finished.
@@ -1096,7 +1096,7 @@ async def test_concurrent_prepare_tables_no_race_condition():
10961096

10971097
# Launch several concurrent create_session calls, each with a unique
10981098
# app_name to avoid IntegrityError on the shared app_states row.
1099-
# Each will call _prepare_tables internally. If the race condition
1099+
# Each will call prepare_tables internally. If the race condition
11001100
# exists, some of these will fail because the "sessions" table doesn't
11011101
# exist yet.
11021102
num_concurrent = 5
@@ -1114,7 +1114,7 @@ async def test_concurrent_prepare_tables_no_race_condition():
11141114
for i, result in enumerate(results):
11151115
assert not isinstance(result, BaseException), (
11161116
f'Concurrent create_session #{i} raised {result!r}; tables were'
1117-
' likely not ready due to the _prepare_tables race condition.'
1117+
' likely not ready due to the prepare_tables race condition.'
11181118
)
11191119

11201120
# All sessions should be retrievable.
@@ -1133,17 +1133,17 @@ async def test_concurrent_prepare_tables_no_race_condition():
11331133
async def test_prepare_tables_serializes_schema_detection_and_creation():
11341134
"""Verifies schema detection and table creation happen atomically under one
11351135
lock, so concurrent callers cannot observe a partially-initialized state.
1136-
After _prepare_tables completes, both _db_schema_version and _tables_created
1136+
After prepare_tables completes, both _db_schema_version and _tables_created
11371137
must be set.
11381138
"""
11391139
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
11401140
try:
11411141
assert not service._tables_created
11421142
assert service._db_schema_version is None
11431143

1144-
await service._prepare_tables()
1144+
await service.prepare_tables()
11451145

1146-
# Both must be set after a single _prepare_tables call.
1146+
# Both must be set after a single prepare_tables call.
11471147
assert service._tables_created
11481148
assert service._db_schema_version is not None
11491149

@@ -1159,17 +1159,17 @@ async def test_prepare_tables_serializes_schema_detection_and_creation():
11591159

11601160
@pytest.mark.asyncio
11611161
async def test_prepare_tables_idempotent_after_creation():
1162-
"""Calling _prepare_tables multiple times is safe and idempotent.
1162+
"""Calling prepare_tables multiple times is safe and idempotent.
11631163
After tables are created, subsequent calls should return immediately via
11641164
the fast path without errors.
11651165
"""
11661166
service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
11671167
try:
1168-
await service._prepare_tables()
1168+
await service.prepare_tables()
11691169
assert service._tables_created
11701170

11711171
# Call again — should be a no-op via the fast path.
1172-
await service._prepare_tables()
1172+
await service.prepare_tables()
11731173
assert service._tables_created
11741174

11751175
# Service should still work.
@@ -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)