Skip to content

Commit 0e57908

Browse files
zouyi100ZOU Yi (BD/SWD-WDE1)
authored andcommitted
fix(sessions): correct timezone handling for PostgreSQL in DatabaseSessionService
`create_session` already stored UTC naive datetimes for both SQLite and PostgreSQL, but the corresponding read path (`get_update_timestamp`) and write path in `append_event` only handled SQLite, causing incorrect POSIX timestamps and potential false stale-session errors on non-UTC hosts. - `get_update_timestamp` / `update_timestamp_tz`: treat PostgreSQL the same as SQLite — attach UTC tzinfo before calling `.timestamp()`. - `to_session`: thread the new `is_postgresql` flag through to `get_update_timestamp`. - `append_event`: use UTC naive datetime for PostgreSQL `update_time` (consistent with `create_session`); use UTC-aware datetime for all other non-SQLite dialects instead of local-time naive. - `get_session` / `list_sessions`: propagate `is_postgresql` to `to_session`.
1 parent b8e8f6b commit 0e57908

File tree

3 files changed

+58
-30
lines changed

3 files changed

+58
-30
lines changed

src/google/adk/sessions/database_session_service.py

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -480,7 +480,7 @@ async def create_session(
480480
storage_app_state.state, storage_user_state.state, session_state
481481
)
482482
session = storage_session.to_session(
483-
state=merged_state, is_sqlite=is_sqlite
483+
state=merged_state, is_sqlite=is_sqlite, is_postgresql=is_postgresql
484484
)
485485
return session
486486

@@ -498,6 +498,8 @@ async def get_session(
498498
# 2. Get all the events based on session id and filtering config
499499
# 3. Convert and return the session
500500
schema = self._get_schema_classes()
501+
is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
502+
is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
501503
async with self._rollback_on_exception_session(
502504
read_only=True
503505
) as sql_session:
@@ -543,9 +545,11 @@ async def get_session(
543545

544546
# Convert storage session to session
545547
events = [e.to_event() for e in reversed(storage_events)]
546-
is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
547548
session = storage_session.to_session(
548-
state=merged_state, events=events, is_sqlite=is_sqlite
549+
state=merged_state,
550+
events=events,
551+
is_sqlite=is_sqlite,
552+
is_postgresql=is_postgresql,
549553
)
550554
return session
551555

@@ -592,12 +596,17 @@ async def list_sessions(
592596

593597
sessions = []
594598
is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
599+
is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
595600
for storage_session in results:
596601
session_state = storage_session.state
597602
user_state = user_states_map.get(storage_session.user_id, {})
598603
merged_state = _merge_state(app_state, user_state, session_state)
599604
sessions.append(
600-
storage_session.to_session(state=merged_state, is_sqlite=is_sqlite)
605+
storage_session.to_session(
606+
state=merged_state,
607+
is_sqlite=is_sqlite,
608+
is_postgresql=is_postgresql,
609+
)
601610
)
602611
return ListSessionsResponse(sessions=sessions)
603612

@@ -633,6 +642,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
633642
# 3. Store the new event.
634643
schema = self._get_schema_classes()
635644
is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
645+
is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
636646
use_row_level_locking = self._supports_row_level_locking()
637647

638648
state_delta = (
@@ -662,7 +672,9 @@ async def append_event(self, session: Session, event: Event) -> Event:
662672
storage_session = storage_session_result.scalars().one_or_none()
663673
if storage_session is None:
664674
raise ValueError(f"Session {session.id} not found.")
665-
storage_update_time = storage_session.get_update_timestamp(is_sqlite)
675+
storage_update_time = storage_session.get_update_timestamp(
676+
is_sqlite, is_postgresql
677+
)
666678
storage_update_marker = storage_session.get_update_marker()
667679

668680
storage_app_state = await _select_required_state(
@@ -728,20 +740,20 @@ async def append_event(self, session: Session, event: Event) -> Event:
728740
storage_session.state | state_deltas["session"]
729741
)
730742

731-
if is_sqlite:
743+
if is_sqlite or is_postgresql:
732744
update_time = datetime.fromtimestamp(
733745
event.timestamp, timezone.utc
734746
).replace(tzinfo=None)
735747
else:
736-
update_time = datetime.fromtimestamp(event.timestamp)
748+
update_time = datetime.fromtimestamp(event.timestamp, timezone.utc)
737749
storage_session.update_time = update_time
738750
sql_session.add(schema.StorageEvent.from_event(session, event))
739751

740752
await sql_session.commit()
741753

742754
# Update timestamp with commit time
743755
session.last_update_time = storage_session.get_update_timestamp(
744-
is_sqlite
756+
is_sqlite, is_postgresql
745757
)
746758
session._storage_update_marker = storage_session.get_update_marker()
747759

src/google/adk/sessions/schemas/v0.py

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -167,19 +167,24 @@ def update_timestamp_tz(self) -> float:
167167
This is a compatibility alias for callers that used the pre-`main` API.
168168
"""
169169
sqlalchemy_session = inspect(self).session
170-
is_sqlite = bool(
171-
sqlalchemy_session
172-
and sqlalchemy_session.bind
173-
and sqlalchemy_session.bind.dialect.name == "sqlite"
170+
dialect_name = (
171+
sqlalchemy_session.bind.dialect.name
172+
if sqlalchemy_session and sqlalchemy_session.bind
173+
else None
174+
)
175+
is_sqlite = dialect_name == "sqlite"
176+
is_postgresql = dialect_name == "postgresql"
177+
return self.get_update_timestamp(
178+
is_sqlite=is_sqlite, is_postgresql=is_postgresql
174179
)
175-
return self.get_update_timestamp(is_sqlite=is_sqlite)
176180

177-
def get_update_timestamp(self, is_sqlite: bool) -> float:
181+
def get_update_timestamp(
182+
self, is_sqlite: bool, is_postgresql: bool = False
183+
) -> float:
178184
"""Returns the time zone aware update timestamp."""
179-
if is_sqlite:
180-
# SQLite does not support timezone. SQLAlchemy returns a naive datetime
181-
# object without timezone information. We need to convert it to UTC
182-
# manually.
185+
if is_sqlite or is_postgresql:
186+
# SQLite and PostgreSQL store naive datetimes as UTC values. We need to
187+
# attach UTC timezone info before converting to a POSIX timestamp.
183188
return self.update_time.replace(tzinfo=timezone.utc).timestamp()
184189
return self.update_time.timestamp()
185190

@@ -195,6 +200,7 @@ def to_session(
195200
state: dict[str, Any] | None = None,
196201
events: list[Event] | None = None,
197202
is_sqlite: bool = False,
203+
is_postgresql: bool = False,
198204
) -> Session:
199205
"""Converts the storage session to a session object."""
200206
if state is None:
@@ -208,7 +214,9 @@ def to_session(
208214
id=self.id,
209215
state=state,
210216
events=events,
211-
last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
217+
last_update_time=self.get_update_timestamp(
218+
is_sqlite=is_sqlite, is_postgresql=is_postgresql
219+
),
212220
)
213221
session._storage_update_marker = self.get_update_marker()
214222
return session

src/google/adk/sessions/schemas/v1.py

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -114,19 +114,24 @@ def update_timestamp_tz(self) -> float:
114114
This is a compatibility alias for callers that used the pre-`main` API.
115115
"""
116116
sqlalchemy_session = inspect(self).session
117-
is_sqlite = bool(
118-
sqlalchemy_session
119-
and sqlalchemy_session.bind
120-
and sqlalchemy_session.bind.dialect.name == "sqlite"
117+
dialect_name = (
118+
sqlalchemy_session.bind.dialect.name
119+
if sqlalchemy_session and sqlalchemy_session.bind
120+
else None
121+
)
122+
is_sqlite = dialect_name == "sqlite"
123+
is_postgresql = dialect_name == "postgresql"
124+
return self.get_update_timestamp(
125+
is_sqlite=is_sqlite, is_postgresql=is_postgresql
121126
)
122-
return self.get_update_timestamp(is_sqlite=is_sqlite)
123127

124-
def get_update_timestamp(self, is_sqlite: bool) -> float:
128+
def get_update_timestamp(
129+
self, is_sqlite: bool, is_postgresql: bool = False
130+
) -> float:
125131
"""Returns the time zone aware update timestamp."""
126-
if is_sqlite:
127-
# SQLite does not support timezone. SQLAlchemy returns a naive datetime
128-
# object without timezone information. We need to convert it to UTC
129-
# manually.
132+
if is_sqlite or is_postgresql:
133+
# SQLite and PostgreSQL store naive datetimes as UTC values. We need to
134+
# attach UTC timezone info before converting to a POSIX timestamp.
130135
return self.update_time.replace(tzinfo=timezone.utc).timestamp()
131136
return self.update_time.timestamp()
132137

@@ -142,6 +147,7 @@ def to_session(
142147
state: dict[str, Any] | None = None,
143148
events: list[Event] | None = None,
144149
is_sqlite: bool = False,
150+
is_postgresql: bool = False,
145151
) -> Session:
146152
"""Converts the storage session to a session object."""
147153
if state is None:
@@ -155,7 +161,9 @@ def to_session(
155161
id=self.id,
156162
state=state,
157163
events=events,
158-
last_update_time=self.get_update_timestamp(is_sqlite=is_sqlite),
164+
last_update_time=self.get_update_timestamp(
165+
is_sqlite=is_sqlite, is_postgresql=is_postgresql
166+
),
159167
)
160168
session._storage_update_marker = self.get_update_marker()
161169
return session

0 commit comments

Comments
 (0)