fix(sessions): strip tzinfo for MariaDB in DatabaseSessionService

MariaDB stores DATETIME/TIMESTAMP values without timezone information,
exactly like MySQL. create_session previously produced a timezone-aware
_storage_update_marker for MariaDB, but the value read back from storage is
naive, so the strict marker comparison in append_event raised a false
"The session has been modified in storage" ValueError on the first
append_event after create_session.

Add MariaDB to the set of dialects whose datetimes are stored as naive, and
consolidate the timezone-naive-dialect decision into a single
_NAIVE_DATETIME_DIALECTS constant and a _uses_naive_datetime() helper used by
create_session, so the strip logic lives in one place. Cloud Spanner is
intentionally excluded because its TIMESTAMP is timezone-aware.

Co-authored-by: Haran Rajkumar <haranrk@google.com>
PiperOrigin-RevId: 937063102

Author: haranrk

src/google/adk/sessions/database_session_service.py

@@ -79,6 +79,18 @@
 _MARIADB_DIALECT = "mariadb"
 _MYSQL_DIALECT = "mysql"
 _POSTGRESQL_DIALECT = "postgresql"
+# Dialects whose DATETIME/TIMESTAMP columns do not retain timezone info, so
+# timezone-aware datetimes must have their tzinfo stripped before storage. This
+# keeps the value written by create_session consistent with the value read back
+# from storage; otherwise the stale-writer marker comparison in append_event
+# raises a false positive on the first append after create_session. Cloud
+# Spanner is intentionally excluded because its TIMESTAMP is timezone-aware.
+_NAIVE_DATETIME_DIALECTS = (
+    _SQLITE_DIALECT,
+    _POSTGRESQL_DIALECT,
+    _MYSQL_DIALECT,
+    _MARIADB_DIALECT,
+)
 # Tuple key order for in-process per-session lock maps:
 # (app_name, user_id, session_id).
 _SessionLockKey: TypeAlias = tuple[str, str, str]
@@ -348,6 +360,14 @@ def _supports_row_level_locking(self) -> bool:
         _POSTGRESQL_DIALECT,
     )
 
+  def _uses_naive_datetime(self) -> bool:
+    """Returns whether the active dialect stores datetimes without timezone info.
+
+    These dialects persist timezone-naive DATETIME/TIMESTAMP values, so
+    timezone-aware datetimes must have their tzinfo stripped before storage.
+    """
+    return self.db_engine.dialect.name in _NAIVE_DATETIME_DIALECTS
+
   @asynccontextmanager
   async def _with_session_lock(
       self, *, app_name: str, user_id: str, session_id: str
@@ -531,8 +551,7 @@ async def create_session(
       now = datetime.fromtimestamp(platform_time.get_time(), tz=timezone.utc)
       is_sqlite = self.db_engine.dialect.name == _SQLITE_DIALECT
       is_postgresql = self.db_engine.dialect.name == _POSTGRESQL_DIALECT
-      is_mysql = self.db_engine.dialect.name == _MYSQL_DIALECT
-      if is_sqlite or is_postgresql or is_mysql:
+      if self._uses_naive_datetime():
         now = now.replace(tzinfo=None)
 
       storage_session = schema.StorageSession(

tests/unittests/sessions/test_session_service.py

@@ -106,29 +106,39 @@ def fake_create_async_engine(_db_url: str, **kwargs):
   assert captured_kwargs.get('pool_pre_ping') is True
 
 
-@pytest.mark.parametrize('dialect_name', ['sqlite', 'postgresql', 'mysql'])
-def test_database_session_service_strips_timezone_for_dialect(dialect_name):
-  """Verifies that timezone-aware datetimes are converted to naive datetimes
-  for SQLite, PostgreSQL, and MySQL.
-
-  These databases store DATETIME/TIMESTAMP WITHOUT TIME ZONE, so keeping
-  tzinfo on the Python datetime causes a mismatch between the marker
-  produced by create_session (with +00:00) and the marker read back by
-  get_session (without +00:00), triggering a false stale-writer error
-  on the first append_event after create_session.
+@pytest.mark.parametrize(
+    'dialect_name', ['sqlite', 'postgresql', 'mysql', 'mariadb']
+)
+def test_database_session_service_uses_naive_datetime_for_dialect(dialect_name):
+  """Verifies dialects that store DATETIME WITHOUT TIME ZONE are treated as naive.
+
+  SQLite, PostgreSQL, MySQL, and MariaDB all store DATETIME/TIMESTAMP WITHOUT
+  TIME ZONE, so create_session must strip tzinfo before storing. Otherwise the
+  marker produced by create_session (with +00:00) mismatches the marker read
+  back from storage (without +00:00), triggering a false stale-writer error on
+  the first append_event after create_session.
+
+  This exercises the production decision (_uses_naive_datetime) directly rather
+  than re-implementing the strip logic, so it actually guards create_session.
   """
-  # Simulate the logic in create_session
-  is_sqlite = dialect_name == 'sqlite'
-  is_postgres = dialect_name == 'postgresql'
-  is_mysql = dialect_name == 'mysql'
+  fake_engine = mock.Mock()
+  fake_engine.dialect.name = dialect_name
+  fake_engine.sync_engine = mock.Mock()
 
-  now = datetime.now(timezone.utc)
-  assert now.tzinfo is not None  # Starts with timezone
+  service = DatabaseSessionService(db_engine=fake_engine)
+
+  assert service._uses_naive_datetime() is True
+
+
+def test_database_session_service_keeps_timezone_for_spanner():
+  """Spanner's TIMESTAMP is timezone-aware, so tzinfo must be preserved."""
+  fake_engine = mock.Mock()
+  fake_engine.dialect.name = 'spanner+spanner'
+  fake_engine.sync_engine = mock.Mock()
 
-  if is_sqlite or is_postgres or is_mysql:
-    now = now.replace(tzinfo=None)
+  service = DatabaseSessionService(db_engine=fake_engine)
 
-  assert now.tzinfo is None
+  assert service._uses_naive_datetime() is False
 
 
 def test_database_session_service_respects_pool_pre_ping_override():