fix(sessions): use INSERT ON CONFLICT DO NOTHING to prevent race in create_session#4955
fix(sessions): use INSERT ON CONFLICT DO NOTHING to prevent race in create_session#4955Lyt060814 wants to merge 2 commits intogoogle:mainfrom
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
|
Hi @Lyt060814 , Thank you for your contribution! It appears you haven't yet signed the Contributor License Agreement (CLA). Please visit https://cla.developers.google.com/ to complete the signing process. Once the CLA is signed, we'll be able to proceed with the review of your PR. Thank you! |
…reate_session When multiple concurrent create_session calls target the same app_name on a fresh database, the SELECT-then-INSERT pattern on app_states and user_states tables causes UniqueViolation errors. All concurrent callers see None from the SELECT and all attempt INSERT. Replace the check-then-insert pattern with dialect-specific INSERT ... ON CONFLICT DO NOTHING (PostgreSQL and SQLite). For unsupported dialects, the original get-then-add fallback is preserved. Added regression tests for concurrent session creation with: - Same app_name, different user_ids (app_states race) - Same app_name AND user_id (user_states race) Fixes google#4954
042df3d to
a99b903
Compare
|
Hi @rohityan, thanks for the heads up! CLA is now signed and the check is passing ✅. I've also rebased onto the latest main to bring the branch up to date. Ready for review whenever you get a chance! |
Description
DatabaseSessionService.create_sessioncontains a SELECT-then-INSERT pattern when initializingapp_statesanduser_statesrows for a newapp_name/user_id. Under concurrent parallel calls (e.g. multiple sub-agents starting simultaneously), all callers seeNonefrom the SELECT and all attempt INSERT, causing aUniqueViolationon the primary key constraint.Fix
Replace the check-then-insert pattern with dialect-specific
INSERT ... ON CONFLICT DO NOTHING:sqlalchemy.dialects.postgresql.insertwithon_conflict_do_nothing(index_elements=[...])sqlalchemy.dialects.sqlite.insertwithon_conflict_do_nothing()The upsert is followed by a regular
SELECTto fetch the (now guaranteed to exist) state rows for subsequent state delta application.Tests
Added 2 regression tests:
test_concurrent_create_session_no_unique_violation: 10 concurrentcreate_sessioncalls with sameapp_name, differentuser_ids (tests theapp_statesrace)test_concurrent_create_session_same_user_no_unique_violation: 10 concurrent calls with sameapp_nameANDuser_id(tests theuser_statesrace)All 80 existing session tests pass.
Files Changed
src/google/adk/sessions/database_session_service.py: Added_ensure_state_rows()method, updatedcreate_session()to use ittests/unittests/sessions/test_session_service.py: Added 2 regression testsFixes #4954