feat: Use new JSON-based schema for newly-created databases

Part 2 of https://github.com/google/adk-python/discussions/3605.

The DatabaseSessionService now checks for the usage of a V1 schema based on the "adk_internal_metadata" table. Table creation and subsequent operations use either the V0 or V1 SQLAlchemy models accordingly. New databases will default to V1.

Migration script and CLI command will be provided in the next change.

Co-authored-by: Liang Wu <wuliang@google.com>
PiperOrigin-RevId: 845443406

src/google/adk/sessions/database_session_service.py

@@ -16,6 +16,7 @@ from __future__ import annotations
 import asyncio
 import copy
 from datetime import datetime
+from datetime import timezone
 import logging
 from typing import Any
 from typing import Optional
@@ -23,12 +24,13 @@ from typing import Optional
 from sqlalchemy import delete
 from sqlalchemy import event
 from sqlalchemy import select
+from sqlalchemy import text
 from sqlalchemy.exc import ArgumentError
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
 from sqlalchemy.ext.asyncio import create_async_engine
-from sqlalchemy.schema import MetaData
+from sqlalchemy.inspection import inspect
 from typing_extensions import override
 from tzlocal import get_localzone
 
@@ -38,11 +40,18 @@ from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
 from .base_session_service import ListSessionsResponse
+from .migration import _schema_check_utils
 from .schemas.v0 import Base as BaseV0
 from .schemas.v0 import StorageAppState as StorageAppStateV0
 from .schemas.v0 import StorageEvent as StorageEventV0
 from .schemas.v0 import StorageSession as StorageSessionV0
 from .schemas.v0 import StorageUserState as StorageUserStateV0
+from .schemas.v1 import Base as BaseV1
+from .schemas.v1 import StorageAppState as StorageAppStateV1
+from .schemas.v1 import StorageEvent as StorageEventV1
+from .schemas.v1 import StorageMetadata
+from .schemas.v1 import StorageSession as StorageSessionV1
+from .schemas.v1 import StorageUserState as StorageUserStateV1
 from .session import Session
 from .state import State
 
@@ -69,6 +78,22 @@ def _merge_state(
   return merged_state
 
 
+class _SchemaClasses:
+  """A helper class to hold schema classes based on version."""
+
+  def __init__(self, version: str):
+    if version == _schema_check_utils.LATEST_SCHEMA_VERSION:
+      self.StorageSession = StorageSessionV1
+      self.StorageAppState = StorageAppStateV1
+      self.StorageUserState = StorageUserStateV1
+      self.StorageEvent = StorageEventV1
+    else:
+      self.StorageSession = StorageSessionV0
+      self.StorageAppState = StorageAppStateV0
+      self.StorageUserState = StorageUserStateV0
+      self.StorageEvent = StorageEventV0
+
+
 class DatabaseSessionService(BaseSessionService):
   """A session service that uses a database for storage."""
 
@@ -101,7 +126,6 @@ class DatabaseSessionService(BaseSessionService):
     logger.info("Local timezone: %s", local_timezone)
 
     self.db_engine: AsyncEngine = db_engine
-    self.metadata: MetaData = MetaData()
 
     # DB session factory method
     self.database_session_factory: async_sessionmaker[
@@ -110,11 +134,48 @@ class DatabaseSessionService(BaseSessionService):
 
     # Flag to indicate if tables are created
     self._tables_created = False
+
     # Lock to ensure thread-safe table creation
     self._table_creation_lock = asyncio.Lock()
 
-  async def _ensure_tables_created(self):
-    """Ensure database tables are created. This is called lazily."""
+    # The current database schema version in use, "None" if not yet checked
+    self._db_schema_version: Optional[str] = None
+
+    # Lock to ensure thread-safe schema version check
+    self._db_schema_lock = asyncio.Lock()
+
+  def _get_schema_classes(self) -> _SchemaClasses:
+    return _SchemaClasses(self._db_schema_version)
+
+  async def _prepare_tables(self):
+    """Ensure database tables are ready for use.
+
+    This method is called lazily before each database operation. It checks the
+    DB schema version to use and creates the tables (including setting the
+    schema version metadata) if needed.
+    """
+    # Check the database schema version and set the _db_schema_version if
+    # needed
+    if self._db_schema_version is not None:
+      return
+
+    async with self._db_schema_lock:
+      # Double-check after acquiring the lock
+      if self._db_schema_version is not None:
+        return
+      try:
+        async with self.db_engine.connect() as conn:
+          self._db_schema_version = await conn.run_sync(
+              _schema_check_utils.get_db_schema_version_from_connection
+          )
+      except Exception:
+        # If inspection fails, assume the latest schema
+        logger.warning(
+            "Failed to inspect database tables, assuming the latest schema."
+        )
+        self._db_schema_version = _schema_check_utils.LATEST_SCHEMA_VERSION
+
+    # Check if tables are created and create them if not
     if self._tables_created:
       return
 
@@ -122,11 +183,37 @@ class DatabaseSessionService(BaseSessionService):
       # Double-check after acquiring the lock
       if not self._tables_created:
         async with self.db_engine.begin() as conn:
-          # Uncomment to recreate DB every time
-          # await conn.run_sync(BaseV0.metadata.drop_all)
-          await conn.run_sync(BaseV0.metadata.create_all)
+          if (
+              self._db_schema_version
+              == _schema_check_utils.LATEST_SCHEMA_VERSION
+          ):
+            # Uncomment to recreate DB every time
+            # await conn.run_sync(BaseV1.metadata.drop_all)
+            logger.debug("Using V1 schema tables...")
+            await conn.run_sync(BaseV1.metadata.create_all)
+          else:
+            # await conn.run_sync(BaseV0.metadata.drop_all)
+            logger.debug("Using V0 schema tables...")
+            await conn.run_sync(BaseV0.metadata.create_all)
         self._tables_created = True
 
+        if self._db_schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION:
+          async with self.database_session_factory() as sql_session:
+            # Check if schema version is set, if not, set it to the latest
+            # version
+            stmt = select(StorageMetadata).where(
+                StorageMetadata.key == _schema_check_utils.SCHEMA_VERSION_KEY
+            )
+            result = await sql_session.execute(stmt)
+            metadata = result.scalars().first()
+            if not metadata:
+              metadata = StorageMetadata(
+                  key=_schema_check_utils.SCHEMA_VERSION_KEY,
+                  value=_schema_check_utils.LATEST_SCHEMA_VERSION,
+              )
+              sql_session.add(metadata)
+              await sql_session.commit()
+
   @override
   async def create_session(
       self,
@@ -141,30 +228,29 @@ class DatabaseSessionService(BaseSessionService):
     # 3. Add the object to the table
     # 4. Build the session object with generated id
     # 5. Return the session
-    await self._ensure_tables_created()
+    await self._prepare_tables()
+    schema = self._get_schema_classes()
     async with self.database_session_factory() as sql_session:
-      StorageSession = StorageSessionV0
-      StorageAppState = StorageAppStateV0
-      StorageUserState = StorageUserStateV0
-
       if session_id and await sql_session.get(
-          StorageSession, (app_name, user_id, session_id)
+          schema.StorageSession, (app_name, user_id, session_id)
       ):
         raise AlreadyExistsError(
             f"Session with id {session_id} already exists."
         )
       # Fetch app and user states from storage
-      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_app_state = await sql_session.get(
+          schema.StorageAppState, (app_name)
+      )
       storage_user_state = await sql_session.get(
-          StorageUserState, (app_name, user_id)
+          schema.StorageUserState, (app_name, user_id)
       )
 
       # Create state tables if not exist
       if not storage_app_state:
-        storage_app_state = StorageAppState(app_name=app_name, state={})
+        storage_app_state = schema.StorageAppState(app_name=app_name, state={})
         sql_session.add(storage_app_state)
       if not storage_user_state:
-        storage_user_state = StorageUserState(
+        storage_user_state = schema.StorageUserState(
             app_name=app_name, user_id=user_id, state={}
         )
         sql_session.add(storage_user_state)
@@ -182,7 +268,7 @@ class DatabaseSessionService(BaseSessionService):
         storage_user_state.state = storage_user_state.state | user_state_delta
 
       # Store the session
-      storage_session = StorageSession(
+      storage_session = schema.StorageSession(
           app_name=app_name,
           user_id=user_id,
           id=session_id,
@@ -209,34 +295,30 @@ class DatabaseSessionService(BaseSessionService):
       session_id: str,
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
-    await self._ensure_tables_created()
+    await self._prepare_tables()
     # 1. Get the storage session entry from session table
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
+    schema = self._get_schema_classes()
     async with self.database_session_factory() as sql_session:
-      StorageSession = StorageSessionV0
-      StorageEvent = StorageEventV0
-      StorageAppState = StorageAppStateV0
-      StorageUserState = StorageUserStateV0
-
       storage_session = await sql_session.get(
-          StorageSession, (app_name, user_id, session_id)
+          schema.StorageSession, (app_name, user_id, session_id)
       )
       if storage_session is None:
         return None
 
       stmt = (
-          select(StorageEvent)
-          .filter(StorageEvent.app_name == app_name)
-          .filter(StorageEvent.session_id == storage_session.id)
-          .filter(StorageEvent.user_id == user_id)
+          select(schema.StorageEvent)
+          .filter(schema.StorageEvent.app_name == app_name)
+          .filter(schema.StorageEvent.session_id == storage_session.id)
+          .filter(schema.StorageEvent.user_id == user_id)
       )
 
       if config and config.after_timestamp:
         after_dt = datetime.fromtimestamp(config.after_timestamp)
-        stmt = stmt.filter(StorageEvent.timestamp >= after_dt)
+        stmt = stmt.filter(schema.StorageEvent.timestamp >= after_dt)
 
-      stmt = stmt.order_by(StorageEvent.timestamp.desc())
+      stmt = stmt.order_by(schema.StorageEvent.timestamp.desc())
 
       if config and config.num_recent_events:
         stmt = stmt.limit(config.num_recent_events)
@@ -245,9 +327,11 @@ class DatabaseSessionService(BaseSessionService):
       storage_events = result.scalars().all()
 
       # Fetch states from storage
-      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_app_state = await sql_session.get(
+          schema.StorageAppState, (app_name)
+      )
       storage_user_state = await sql_session.get(
-          StorageUserState, (app_name, user_id)
+          schema.StorageUserState, (app_name, user_id)
       )
 
       app_state = storage_app_state.state if storage_app_state else {}
@@ -266,34 +350,35 @@ class DatabaseSessionService(BaseSessionService):
   async def list_sessions(
       self, *, app_name: str, user_id: Optional[str] = None
   ) -> ListSessionsResponse:
-    await self._ensure_tables_created()
+    await self._prepare_tables()
+    schema = self._get_schema_classes()
     async with self.database_session_factory() as sql_session:
-      StorageSession = StorageSessionV0
-      StorageAppState = StorageAppStateV0
-      StorageUserState = StorageUserStateV0
-
-      stmt = select(StorageSession).filter(StorageSession.app_name == app_name)
+      stmt = select(schema.StorageSession).filter(
+          schema.StorageSession.app_name == app_name
+      )
       if user_id is not None:
-        stmt = stmt.filter(StorageSession.user_id == user_id)
+        stmt = stmt.filter(schema.StorageSession.user_id == user_id)
 
       result = await sql_session.execute(stmt)
       results = result.scalars().all()
 
       # Fetch app state from storage
-      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_app_state = await sql_session.get(
+          schema.StorageAppState, (app_name)
+      )
       app_state = storage_app_state.state if storage_app_state else {}
 
       # Fetch user state(s) from storage
       user_states_map = {}
       if user_id is not None:
         storage_user_state = await sql_session.get(
-            StorageUserState, (app_name, user_id)
+            schema.StorageUserState, (app_name, user_id)
         )
         if storage_user_state:
           user_states_map[user_id] = storage_user_state.state
       else:
-        user_state_stmt = select(StorageUserState).filter(
-            StorageUserState.app_name == app_name
+        user_state_stmt = select(schema.StorageUserState).filter(
+            schema.StorageUserState.app_name == app_name
         )
         user_state_result = await sql_session.execute(user_state_stmt)
         all_user_states_for_app = user_state_result.scalars().all()
@@ -312,21 +397,20 @@ class DatabaseSessionService(BaseSessionService):
   async def delete_session(
       self, app_name: str, user_id: str, session_id: str
   ) -> None:
-    await self._ensure_tables_created()
+    await self._prepare_tables()
+    schema = self._get_schema_classes()
     async with self.database_session_factory() as sql_session:
-      StorageSession = StorageSessionV0
-
-      stmt = delete(StorageSession).where(
-          StorageSession.app_name == app_name,
-          StorageSession.user_id == user_id,
-          StorageSession.id == session_id,
+      stmt = delete(schema.StorageSession).where(
+          schema.StorageSession.app_name == app_name,
+          schema.StorageSession.user_id == user_id,
+          schema.StorageSession.id == session_id,
       )
       await sql_session.execute(stmt)
       await sql_session.commit()
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
-    await self._ensure_tables_created()
+    await self._prepare_tables()
     if event.partial:
       return event
 
@@ -336,14 +420,10 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Check if timestamp is stale
     # 2. Update session attributes based on event config
     # 3. Store event to table
+    schema = self._get_schema_classes()
     async with self.database_session_factory() as sql_session:
-      StorageSession = StorageSessionV0
-      StorageEvent = StorageEventV0
-      StorageAppState = StorageAppStateV0
-      StorageUserState = StorageUserStateV0
-
       storage_session = await sql_session.get(
-          StorageSession, (session.app_name, session.user_id, session.id)
+          schema.StorageSession, (session.app_name, session.user_id, session.id)
       )
 
       if storage_session.update_timestamp_tz > session.last_update_time:
@@ -357,10 +437,10 @@ class DatabaseSessionService(BaseSessionService):
 
       # Fetch states from storage
       storage_app_state = await sql_session.get(
-          StorageAppState, (session.app_name)
+          schema.StorageAppState, (session.app_name)
       )
       storage_user_state = await sql_session.get(
-          StorageUserState, (session.app_name, session.user_id)
+          schema.StorageUserState, (session.app_name, session.user_id)
       )
 
       # Extract state delta
@@ -380,11 +460,13 @@ class DatabaseSessionService(BaseSessionService):
           storage_session.state = storage_session.state | session_state_delta
 
       if storage_session._dialect_name == "sqlite":
-        update_time = datetime.utcfromtimestamp(event.timestamp)
+        update_time = datetime.fromtimestamp(
+            event.timestamp, timezone.utc
+        ).replace(tzinfo=None)
       else:
         update_time = datetime.fromtimestamp(event.timestamp)
       storage_session.update_time = update_time
-      sql_session.add(StorageEvent.from_event(session, event))
+      sql_session.add(schema.StorageEvent.from_event(session, event))
 
       await sql_session.commit()
       await sql_session.refresh(storage_session)

src/google/adk/sessions/migration/_schema_check_utils.py

@@ -0,0 +1,73 @@
+# Copyright 2025 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.
+"""Database schema version check utility."""
+from __future__ import annotations
+
+import logging
+
+from sqlalchemy import inspect
+from sqlalchemy import text
+
+logger = logging.getLogger("google_adk." + __name__)
+
+SCHEMA_VERSION_KEY = "schema_version"
+SCHEMA_VERSION_0_PICKLE = "0"
+SCHEMA_VERSION_1_JSON = "1"
+LATEST_SCHEMA_VERSION = SCHEMA_VERSION_1_JSON
+
+
+def _get_schema_version_impl(inspector, connection) -> str:
+  """Gets DB schema version using inspector and connection."""
+  if inspector.has_table("adk_internal_metadata"):
+    try:
+      result = connection.execute(
+          text("SELECT value FROM adk_internal_metadata WHERE key = :key"),
+          {"key": SCHEMA_VERSION_KEY},
+      ).fetchone()
+      if result:
+        return result[0]
+      else:
+        return LATEST_SCHEMA_VERSION
+    except Exception as e:
+      logger.warning(
+          "Failed to query schema version from adk_internal_metadata,"
+          " assuming the latest schema: %s.",
+          e,
+      )
+      return LATEST_SCHEMA_VERSION
+  # Metadata table doesn't exist, check for v0 schema.
+  # V0 schema has an 'events' table with an 'actions' column.
+  if inspector.has_table("events"):
+    try:
+      cols = {c["name"] for c in inspector.get_columns("events")}
+      if "actions" in cols and "event_data" not in cols:
+        logger.warning(
+            "The database is using the legacy v0 schema, which uses Pickle to"
+            " serialize event actions. The v0 schema will not be supported"
+            " going forward and will be deprecated in a few rollouts. Please"
+            " migrate to the v1 schema which uses JSON serialization for event"
+            " data. The migration command and script will be provided soon."
+        )
+        return SCHEMA_VERSION_0_PICKLE
+    except Exception as e:
+      logger.warning("Failed to inspect 'events' table columns: %s", e)
+      return LATEST_SCHEMA_VERSION
+  # New database, assume the latest schema.
+  return LATEST_SCHEMA_VERSION
+
+
+def get_db_schema_version_from_connection(connection) -> str:
+  """Gets DB schema version from a DB connection."""
+  inspector = inspect(connection)
+  return _get_schema_version_impl(inspector, connection)

tests/unittests/sessions/migration/test_database_schema.py

@@ -0,0 +1,162 @@
+# Copyright 2025 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.
+
+from google.adk.sessions.database_session_service import DatabaseSessionService
+from google.adk.sessions.migration import _schema_check_utils
+from google.adk.sessions.schemas import v0
+import pytest
+from sqlalchemy import inspect
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+
+
+async def create_v0_db(db_path):
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+  engine = create_async_engine(db_url)
+  async with engine.begin() as conn:
+    await conn.run_sync(v0.Base.metadata.create_all)
+  await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_new_db_uses_latest_schema(tmp_path):
+  db_path = tmp_path / 'new_db.db'
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+  session_service = DatabaseSessionService(db_url)
+  assert session_service._db_schema_version is None
+  await session_service.create_session(app_name='my_app', user_id='test_user')
+  assert (
+      session_service._db_schema_version
+      == _schema_check_utils.LATEST_SCHEMA_VERSION
+  )
+
+  # Verify metadata table
+  engine = create_async_engine(db_url)
+  async with engine.connect() as conn:
+    has_metadata_table = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
+    )
+    assert has_metadata_table
+    schema_version = await conn.run_sync(
+        lambda sync_conn: sync_conn.execute(
+            text('SELECT value FROM adk_internal_metadata WHERE key = :key'),
+            {'key': _schema_check_utils.SCHEMA_VERSION_KEY},
+        ).scalar_one_or_none()
+    )
+    assert schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION
+
+    # Verify events table columns for v1
+    event_cols = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).get_columns('events')
+    )
+    event_col_names = {c['name'] for c in event_cols}
+    assert 'event_data' in event_col_names
+    assert 'actions' not in event_col_names
+  await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_existing_v0_db_uses_v0_schema(tmp_path):
+  db_path = tmp_path / 'v0_db.db'
+  await create_v0_db(db_path)
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+  session_service = DatabaseSessionService(db_url)
+
+  assert session_service._db_schema_version is None
+  await session_service.create_session(
+      app_name='my_app', user_id='test_user', session_id='s1'
+  )
+  assert (
+      session_service._db_schema_version
+      == _schema_check_utils.SCHEMA_VERSION_0_PICKLE
+  )
+
+  session = await session_service.get_session(
+      app_name='my_app', user_id='test_user', session_id='s1'
+  )
+  assert session.id == 's1'
+
+  # Verify schema tables
+  engine = create_async_engine(db_url)
+  async with engine.connect() as conn:
+    has_metadata_table = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
+    )
+    assert not has_metadata_table
+
+    # Verify events table columns for v0
+    event_cols = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).get_columns('events')
+    )
+    event_col_names = {c['name'] for c in event_cols}
+    assert 'event_data' not in event_col_names
+    assert 'actions' in event_col_names
+  await engine.dispose()
+
+
+@pytest.mark.asyncio
+async def test_existing_latest_db_uses_latest_schema(tmp_path):
+  db_path = tmp_path / 'new_db.db'
+  db_url = f'sqlite+aiosqlite:///{db_path}'
+
+  # Create session service which creates db with latest schema
+  session_service1 = DatabaseSessionService(db_url)
+  await session_service1.create_session(
+      app_name='my_app', user_id='test_user', session_id='s1'
+  )
+  assert (
+      session_service1._db_schema_version
+      == _schema_check_utils.LATEST_SCHEMA_VERSION
+  )
+
+  # Create another session service on same db and check it detects latest schema
+  session_service2 = DatabaseSessionService(db_url)
+  await session_service2.create_session(
+      app_name='my_app', user_id='test_user2', session_id='s2'
+  )
+  assert (
+      session_service2._db_schema_version
+      == _schema_check_utils.LATEST_SCHEMA_VERSION
+  )
+  s2 = await session_service2.get_session(
+      app_name='my_app', user_id='test_user2', session_id='s2'
+  )
+  assert s2.id == 's2'
+
+  s1 = await session_service2.get_session(
+      app_name='my_app', user_id='test_user', session_id='s1'
+  )
+  assert s1.id == 's1'
+
+  list_sessions_response = await session_service2.list_sessions(
+      app_name='my_app'
+  )
+  assert len(list_sessions_response.sessions) == 2
+
+  # Verify schema tables
+  engine = create_async_engine(db_url)
+  async with engine.connect() as conn:
+    has_metadata_table = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
+    )
+    assert has_metadata_table
+
+    # Verify events table columns for v1
+    event_cols = await conn.run_sync(
+        lambda sync_conn: inspect(sync_conn).get_columns('events')
+    )
+    event_col_names = {c['name'] for c in event_cols}
+    assert 'event_data' in event_col_names
+    assert 'actions' not in event_col_names
+  await engine.dispose()