fix(sessions): Further fixes for DatabaseSessionService

DeanChensj · copybara-github · commit f0ec997bc012 · 2026-06-16T23:58:41.000-07:00
- Fix timezone inconsistency in append_event where it used local naive time for Postgres (now uses UTC naive).
- Fix potential MissingGreenlet in create_session by generating UUID in Python and calling to_session before commit.
- Add regression test for create_session.

Co-authored-by: Shangjie Chen &lt;deanchen@google.com&gt;
PiperOrigin-RevId: 933531685
diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py
@@ -26,6 +26,7 @@
 from typing import TypeVar
 
 from google.adk.platform import time as platform_time
+from google.adk.platform import uuid as platform_uuid
 
 try:
   from sqlalchemy import delete
@@ -434,9 +435,12 @@ async def create_session(
     # 4. Build the session object with generated id
     # 5. Return the session
     await self._prepare_tables()
+    has_user_provided_id = session_id is not None
+    if session_id is None:
+      session_id = platform_uuid.new_uuid()
     schema = self._get_schema_classes()
     async with self._rollback_on_exception_session() as sql_session:
-      if session_id and await sql_session.get(
+      if has_user_provided_id and await sql_session.get(
           schema.StorageSession, (app_name, user_id, session_id)
       ):
         raise AlreadyExistsError(
@@ -484,15 +488,17 @@ async def create_session(
           update_time=now,
       )
       sql_session.add(storage_session)
-      await sql_session.commit()
 
       # Merge states for response
       merged_state = _merge_state(
           storage_app_state.state, storage_user_state.state, session_state
       )
+      # Call to_session before commit to avoid post-commit lazy-load.
+      await sql_session.flush()
       session = storage_session.to_session(
-          state=merged_state, is_sqlite=is_sqlite
+          state=merged_state, is_sqlite=is_sqlite, is_postgresql=is_postgresql
       )
+      await sql_session.commit()
     return session
 
   @override
@@ -555,8 +561,12 @@ async def get_session(
       # Convert storage session to session
       events = [e.to_event() for e in reversed(storage_events)]
       is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
+      is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
       session = storage_session.to_session(
-          state=merged_state, events=events, is_sqlite=is_sqlite
+          state=merged_state,
+          events=events,
+          is_sqlite=is_sqlite,
+          is_postgresql=is_postgresql,
       )
     return session
 
@@ -603,12 +613,17 @@ async def list_sessions(
 
       sessions = []
       is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
+      is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
       for storage_session in results:
         session_state = storage_session.state
         user_state = user_states_map.get(storage_session.user_id, {})
         merged_state = _merge_state(app_state, user_state, session_state)
         sessions.append(
-            storage_session.to_session(state=merged_state, is_sqlite=is_sqlite)
+            storage_session.to_session(
+                state=merged_state,
+                is_sqlite=is_sqlite,
+                is_postgresql=is_postgresql,
+            )
         )
       return ListSessionsResponse(sessions=sessions)
 
@@ -660,6 +675,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # 3. Store the new event.
     schema = self._get_schema_classes()
     is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
+    is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
     use_row_level_locking = self._supports_row_level_locking()
 
     state_delta = event.actions.state_delta if event.actions.state_delta else {}
@@ -685,7 +701,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
         storage_session = storage_session_result.scalars().one_or_none()
         if storage_session is None:
           raise ValueError(f"Session {session.id} not found.")
-        storage_update_time = storage_session.get_update_timestamp(is_sqlite)
+        storage_update_time = storage_session.get_update_timestamp(
+            is_sqlite=is_sqlite, is_postgresql=is_postgresql
+        )
         storage_update_marker = storage_session.get_update_marker()
 
         storage_app_state = await _select_required_state(
@@ -751,7 +769,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
               storage_session.state | state_deltas["session"]
           )
 
-        if is_sqlite:
+        is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
+        if is_sqlite or is_postgresql:
           update_time = datetime.fromtimestamp(
               event.timestamp, timezone.utc
           ).replace(tzinfo=None)
@@ -763,7 +782,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
         # Read revision fields before commit. Post-commit ORM attribute access
         # can lazy-load expired columns and trigger MissingGreenlet with asyncpg
         # when pool_pre_ping is enabled.
-        last_update_time = storage_session.get_update_timestamp(is_sqlite)
+        last_update_time = storage_session.get_update_timestamp(
+            is_sqlite=is_sqlite, is_postgresql=is_postgresql
+        )
         storage_update_marker = storage_session.get_update_marker()
         await sql_session.commit()
 
diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py
@@ -172,12 +172,22 @@ def update_timestamp_tz(self) -> float:
         and sqlalchemy_session.bind
         and sqlalchemy_session.bind.dialect.name == "sqlite"
     )
-    return self.get_update_timestamp(is_sqlite=is_sqlite)
+    is_postgresql = bool(
+        sqlalchemy_session
+        and sqlalchemy_session.bind
+        and sqlalchemy_session.bind.dialect.name == "postgresql"
+    )
+    return self.get_update_timestamp(
+        is_sqlite=is_sqlite, is_postgresql=is_postgresql
+    )
 
-  def get_update_timestamp(self, is_sqlite: bool) -> float:
+  def get_update_timestamp(
+      self, is_sqlite: bool = False, is_postgresql: bool = False
+  ) -> float:
     """Returns the time zone aware update timestamp."""
-    if is_sqlite:
-      # SQLite does not support timezone. SQLAlchemy returns a naive datetime
+    del is_sqlite, is_postgresql  # Unused.
+    if self.update_time.tzinfo is None:
+      # SQLite and PostgreSQL do not support timezone. SQLAlchemy returns a naive datetime
       # object without timezone information. We need to convert it to UTC
       # manually.
       return self.update_time.replace(tzinfo=timezone.utc).timestamp()
@@ -195,6 +205,7 @@ def to_session(
       state: dict[str, Any] | None = None,
       events: list[Event] | None = None,
       is_sqlite: bool = False,
+      is_postgresql: bool = False,
   ) -> Session:
     """Converts the storage session to a session object."""
     if state is None:
@@ -208,7 +219,9 @@ def to_session(
         id=self.id,
         state=state,
         events=events,
-        last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
+        last_update_time=self.get_update_timestamp(
+            is_sqlite=is_sqlite, is_postgresql=is_postgresql
+        ),
     )
     session._storage_update_marker = self.get_update_marker()
     return session
diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py
@@ -119,12 +119,22 @@ def update_timestamp_tz(self) -> float:
         and sqlalchemy_session.bind
         and sqlalchemy_session.bind.dialect.name == "sqlite"
     )
-    return self.get_update_timestamp(is_sqlite=is_sqlite)
+    is_postgresql = bool(
+        sqlalchemy_session
+        and sqlalchemy_session.bind
+        and sqlalchemy_session.bind.dialect.name == "postgresql"
+    )
+    return self.get_update_timestamp(
+        is_sqlite=is_sqlite, is_postgresql=is_postgresql
+    )
 
-  def get_update_timestamp(self, is_sqlite: bool) -> float:
+  def get_update_timestamp(
+      self, is_sqlite: bool = False, is_postgresql: bool = False
+  ) -> float:
     """Returns the time zone aware update timestamp."""
-    if is_sqlite:
-      # SQLite does not support timezone. SQLAlchemy returns a naive datetime
+    del is_sqlite, is_postgresql  # Unused.
+    if self.update_time.tzinfo is None:
+      # SQLite and PostgreSQL do not support timezone. SQLAlchemy returns a naive datetime
       # object without timezone information. We need to convert it to UTC
       # manually.
       return self.update_time.replace(tzinfo=timezone.utc).timestamp()
@@ -142,6 +152,7 @@ def to_session(
       state: dict[str, Any] | None = None,
       events: list[Event] | None = None,
       is_sqlite: bool = False,
+      is_postgresql: bool = False,
   ) -> Session:
     """Converts the storage session to a session object."""
     if state is None:
@@ -155,7 +166,9 @@ def to_session(
         id=self.id,
         state=state,
         events=events,
-        last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
+        last_update_time=self.get_update_timestamp(
+            is_sqlite=is_sqlite, is_postgresql=is_postgresql
+        ),
     )
     session._storage_update_marker = self.get_update_marker()
     return session
diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py
@@ -1335,6 +1335,56 @@ def _spy_factory():
     await service.close()
 
 
+@pytest.mark.asyncio
+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()
+  schema = service._get_schema_classes()
+  original_get_update_timestamp = schema.StorageSession.get_update_timestamp
+  original_get_update_marker = schema.StorageSession.get_update_marker
+  revision_read_state = {'committed': False, 'post_commit_reads': 0}
+
+  def _track_revision_read(original):
+    def wrapper(self, *args, **kwargs):
+      if revision_read_state['committed']:
+        revision_read_state['post_commit_reads'] += 1
+      return original(self, *args, **kwargs)
+
+    return wrapper
+
+  schema.StorageSession.get_update_timestamp = _track_revision_read(
+      original_get_update_timestamp
+  )
+  schema.StorageSession.get_update_marker = _track_revision_read(
+      original_get_update_marker
+  )
+
+  try:
+    original_factory = service.database_session_factory
+
+    def _spy_factory():
+      return _CommitOrderSpySession(
+          original_factory(),
+          on_committed=lambda: revision_read_state.update({'committed': True}),
+      )
+
+    service.database_session_factory = _spy_factory
+
+    session = await service.create_session(
+        app_name='app', user_id='user', session_id='s1'
+    )
+
+    assert revision_read_state['post_commit_reads'] == 0
+    assert session.last_update_time is not None
+    assert session._storage_update_marker is not None
+  finally:
+    schema.StorageSession.get_update_timestamp = original_get_update_timestamp
+    schema.StorageSession.get_update_marker = original_get_update_marker
+
+    await service.close()
+
+
 @pytest.mark.asyncio
 async def test_delete_session_calls_rollback_on_commit_failure():
   """Verifies that a commit failure during delete_session triggers an explicit
