From 1c57f8fe7c8c0e765d59b8b4fc966f51a5b3a3ff Mon Sep 17 00:00:00 2001 From: tranminhquang Date: Fri, 6 Feb 2026 11:35:09 +0700 Subject: [PATCH 1/4] fix: use TIMESTAMP WITH TIME ZONE for PostgreSQL in PreciseTimestamp The `PreciseTimestamp` TypeDecorator uses `DateTime` (without timezone) as its default implementation. When used with PostgreSQL + asyncpg, this causes the SQL query to cast parameters as `TIMESTAMP WITHOUT TIME ZONE`, which conflicts with timezone-aware datetime objects (those with `tzinfo=UTC`) that ADK now creates internally (since v1.24.0). This results in the asyncpg error: "can't subtract offset-naive and offset-aware datetimes" The fix adds a PostgreSQL-specific dialect implementation that uses `postgresql.TIMESTAMP(timezone=True)`, consistent with how MySQL already has its own dialect-specific handling (`mysql.DATETIME(fsp=6)`). Fixes #1848 --- src/google/adk/sessions/schemas/shared.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/google/adk/sessions/schemas/shared.py b/src/google/adk/sessions/schemas/shared.py index 25d4ea9e95..64641591ea 100644 --- a/src/google/adk/sessions/schemas/shared.py +++ b/src/google/adk/sessions/schemas/shared.py @@ -64,4 +64,6 @@ class PreciseTimestamp(TypeDecorator): def load_dialect_impl(self, dialect): if dialect.name == "mysql": return dialect.type_descriptor(mysql.DATETIME(fsp=6)) + elif dialect.name == "postgresql": + return dialect.type_descriptor(postgresql.TIMESTAMP(timezone=True)) return self.impl From ba39e556d2a0ef86dac897caf8490cced7cdc368 Mon Sep 17 00:00:00 2001 From: tranminhquang Date: Fri, 6 Feb 2026 11:42:48 +0700 Subject: [PATCH 2/4] feat: add PostgreSQL TIMESTAMPTZ migration script Adds a migration script that converts TIMESTAMP WITHOUT TIME ZONE columns to TIMESTAMP WITH TIME ZONE for existing PostgreSQL databases. Usage: python -m google.adk.sessions.migration.migrate_postgresql_timestamptz \ --db_url postgresql+asyncpg://user:pass@host:port/dbname The script checks each ADK table column, skips non-PostgreSQL databases and columns already using TIMESTAMPTZ, and migrates only what's needed. --- .../migrate_postgresql_timestamptz.py | 150 ++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 src/google/adk/sessions/migration/migrate_postgresql_timestamptz.py diff --git a/src/google/adk/sessions/migration/migrate_postgresql_timestamptz.py b/src/google/adk/sessions/migration/migrate_postgresql_timestamptz.py new file mode 100644 index 0000000000..d68e1d0b0a --- /dev/null +++ b/src/google/adk/sessions/migration/migrate_postgresql_timestamptz.py @@ -0,0 +1,150 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Migration script to convert TIMESTAMP to TIMESTAMPTZ for PostgreSQL. + +Starting from ADK v1.24.0, DatabaseSessionService creates timezone-aware +datetime objects (with tzinfo=UTC). When using PostgreSQL with asyncpg, +this causes a conflict if existing timestamp columns are defined as +TIMESTAMP WITHOUT TIME ZONE, resulting in: + + asyncpg.exceptions.DataError: can't subtract offset-naive and + offset-aware datetimes + +This migration alters all timestamp columns in ADK tables to use +TIMESTAMP WITH TIME ZONE. It is safe to run on existing data as +PostgreSQL will interpret existing naive timestamps as being in the +server's timezone (typically UTC). + +Usage: + python -m google.adk.sessions.migration.migrate_postgresql_timestamptz \ + --db_url postgresql+asyncpg://user:pass@host:port/dbname +""" + +from __future__ import annotations + +import argparse +import logging +import sys + +from sqlalchemy import create_engine +from sqlalchemy import text + +from . import _schema_check_utils + +logger = logging.getLogger("google_adk." + __name__) + +# Columns to migrate: (table_name, column_name) +_TIMESTAMP_COLUMNS = [ + ("sessions", "create_time"), + ("sessions", "update_time"), + ("events", "timestamp"), + ("app_states", "update_time"), + ("user_states", "update_time"), +] + + +def migrate(db_url: str) -> None: + """Migrates TIMESTAMP columns to TIMESTAMP WITH TIME ZONE for PostgreSQL. + + Args: + db_url: The database URL (sync or async format). + """ + sync_url = _schema_check_utils.to_sync_url(db_url) + engine = create_engine(sync_url) + + try: + with engine.begin() as conn: + # Only run on PostgreSQL + if engine.dialect.name != "postgresql": + logger.info( + "Skipping TIMESTAMPTZ migration: not a PostgreSQL database" + " (dialect=%s).", + engine.dialect.name, + ) + return + + migrated = 0 + for table_name, column_name in _TIMESTAMP_COLUMNS: + # Check if table exists + result = conn.execute( + text( + "SELECT data_type FROM information_schema.columns " + "WHERE table_schema = 'public' " + "AND table_name = :table_name " + "AND column_name = :column_name" + ), + {"table_name": table_name, "column_name": column_name}, + ).fetchone() + + if result is None: + logger.debug( + "Skipping %s.%s: column not found.", table_name, column_name + ) + continue + + if result[0] == "timestamp with time zone": + logger.debug( + "Skipping %s.%s: already TIMESTAMP WITH TIME ZONE.", + table_name, + column_name, + ) + continue + + logger.info( + "Migrating %s.%s from %s to TIMESTAMP WITH TIME ZONE.", + table_name, + column_name, + result[0], + ) + conn.execute( + text( + f"ALTER TABLE {table_name} " + f"ALTER COLUMN {column_name} " + f"TYPE TIMESTAMP WITH TIME ZONE" + ) + ) + migrated += 1 + + if migrated > 0: + logger.info( + "Successfully migrated %d column(s) to TIMESTAMP WITH TIME ZONE.", + migrated, + ) + else: + logger.info("No columns needed migration.") + + finally: + engine.dispose() + + +def main(): + parser = argparse.ArgumentParser( + description=( + "Migrate PostgreSQL TIMESTAMP columns to TIMESTAMP WITH TIME ZONE" + " for ADK DatabaseSessionService." + ) + ) + parser.add_argument( + "--db_url", + required=True, + help="Database URL (e.g., postgresql+asyncpg://user:pass@host:port/db)", + ) + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + migrate(args.db_url) + + +if __name__ == "__main__": + main() From 001962a49170322bd86262847b9b997dd369087d Mon Sep 17 00:00:00 2001 From: tranminhquang Date: Fri, 6 Feb 2026 15:19:47 +0700 Subject: [PATCH 3/4] test: add unit tests for PreciseTimestamp dialect-specific type resolution Verifies that PreciseTimestamp returns the correct SQLAlchemy type for each database dialect: - PostgreSQL: TIMESTAMP(timezone=True) to prevent asyncpg offset-naive vs offset-aware datetime errors - MySQL: DATETIME(fsp=6) for microsecond precision - SQLite: default DateTime fallback Includes regression test proving the fix catches the bug: without the PostgreSQL branch, load_dialect_impl returns plain DateTime (timezone=False), which causes asyncpg DataError. --- .../sessions/test_precise_timestamp.py | 69 +++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 tests/unittests/sessions/test_precise_timestamp.py diff --git a/tests/unittests/sessions/test_precise_timestamp.py b/tests/unittests/sessions/test_precise_timestamp.py new file mode 100644 index 0000000000..206d14624e --- /dev/null +++ b/tests/unittests/sessions/test_precise_timestamp.py @@ -0,0 +1,69 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for PreciseTimestamp dialect-specific type resolution. + +These tests verify that PreciseTimestamp maps to the correct database-specific +column types, particularly ensuring PostgreSQL uses TIMESTAMP WITH TIME ZONE +to prevent asyncpg 'offset-naive and offset-aware datetimes' errors. +""" + +from google.adk.sessions.schemas.shared import PreciseTimestamp +from sqlalchemy.dialects import mysql +from sqlalchemy.dialects import postgresql +from sqlalchemy.dialects import sqlite +from sqlalchemy.types import DateTime + + +class TestPreciseTimestampDialectImpl: + """Tests that PreciseTimestamp.load_dialect_impl returns the correct type.""" + + def test_postgresql_returns_timestamp_with_timezone(self): + """PostgreSQL must use TIMESTAMP WITH TIME ZONE to accept timezone-aware + datetimes from asyncpg without raising DataError.""" + ts = PreciseTimestamp() + dialect = postgresql.dialect() + impl = ts.load_dialect_impl(dialect) + assert isinstance(impl, postgresql.TIMESTAMP) + assert impl.timezone is True + + def test_postgresql_not_timestamp_without_timezone(self): + """Regression test: PostgreSQL must NOT use TIMESTAMP WITHOUT TIME ZONE. + + Without timezone=True, asyncpg raises: + DataError: can't subtract offset-naive and offset-aware datetimes + when inserting datetime objects with tzinfo=UTC. + """ + ts = PreciseTimestamp() + dialect = postgresql.dialect() + impl = ts.load_dialect_impl(dialect) + # Ensure it's not the default TIMESTAMP (which has timezone=False) + naive_timestamp = postgresql.TIMESTAMP() + assert naive_timestamp.timezone is False # baseline: default is False + assert impl.timezone is not naive_timestamp.timezone + + def test_mysql_returns_datetime_with_fsp6(self): + """MySQL must use DATETIME(fsp=6) for microsecond precision.""" + ts = PreciseTimestamp() + dialect = mysql.dialect() + impl = ts.load_dialect_impl(dialect) + assert isinstance(impl, mysql.DATETIME) + assert impl.fsp == 6 + + def test_sqlite_returns_default_datetime(self): + """SQLite falls back to the default DateTime implementation.""" + ts = PreciseTimestamp() + dialect = sqlite.dialect() + impl = ts.load_dialect_impl(dialect) + assert isinstance(impl, DateTime) From d93865767064ecfcc9ee380fe0ee9665ebce07fe Mon Sep 17 00:00:00 2001 From: vietnamesekid Date: Fri, 6 Feb 2026 16:16:11 +0700 Subject: [PATCH 4/4] fix: use explicit UTC datetimes in append_event and get_session for non-SQLite dialects MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix timezone-naive datetime creation in database_session_service.py that causes incorrect timestamps on servers not running in UTC. Changes: - append_event: use datetime.fromtimestamp(ts, timezone.utc) for non-SQLite dialects when setting session update_time - get_session: use datetime.fromtimestamp(ts, timezone.utc) for non-SQLite dialects when filtering events by after_timestamp Why only database_session_service.py and not schemas/v0.py or v1.py: The schema-level from_event() methods (StorageEvent.from_event) are dialect-agnostic — they have no access to the database engine to know which dialect is in use. Changing them to produce UTC-aware datetimes breaks SQLite roundtrip: SQLite stores datetimes as naive strings, so an aware datetime gets its tzinfo stripped on write, but .timestamp() on the naive readback interprets it as local time, producing wrong values on non-UTC servers (e.g. UTC+7: epoch 3 → stored as 1970-01-01 00:00:03 → read back as naive → .timestamp() = -28797). The service layer (database_session_service.py) is the correct place for this fix because it has access to self.db_engine.dialect.name and can branch on SQLite vs PostgreSQL/MySQL accordingly. --- src/google/adk/sessions/database_session_service.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 18dd999a94..1b5a396eb9 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -344,7 +344,12 @@ async def get_session( ) if config and config.after_timestamp: - after_dt = datetime.fromtimestamp(config.after_timestamp) + if self.db_engine.dialect.name == "sqlite": + after_dt = datetime.fromtimestamp(config.after_timestamp) + else: + after_dt = datetime.fromtimestamp( + config.after_timestamp, timezone.utc + ) stmt = stmt.filter(schema.StorageEvent.timestamp >= after_dt) stmt = stmt.order_by(schema.StorageEvent.timestamp.desc()) @@ -513,7 +518,7 @@ async def append_event(self, session: Session, event: Event) -> Event: event.timestamp, timezone.utc ).replace(tzinfo=None) else: - update_time = datetime.fromtimestamp(event.timestamp) + update_time = datetime.fromtimestamp(event.timestamp, timezone.utc) storage_session.update_time = update_time sql_session.add(schema.StorageEvent.from_event(session, event))