feat: Add database schema migration command and script

Final part of #3605.

This change introduces:
-   A new `adk migrate session` CLI command to run database schema upgrades.
-   A migration script to upgrade from the old Pickle-based session schema (v0) to the new JSON-based schema (v1).
-   A migration runner that orchestrates the upgrade process, handling sequential migrations and using temporary SQLite databases for intermediate steps if needed.
-   Unit tests for the v0 to v1 migration.

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

Author: wuliang229

src/google/adk/cli/cli_tools_click.py

@@ -36,6 +36,7 @@
 from . import cli_deploy
 from .. import version
 from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
+from ..sessions.migration import migration_runner
 from .cli import run_cli
 from .fast_api import get_fast_api_app
 from .utils import envs
@@ -1507,6 +1508,47 @@ def cli_deploy_cloud_run(
     click.secho(f"Deploy failed: {e}", fg="red", err=True)
 
 
+@main.group()
+def migrate():
+  """ADK migration commands."""
+  pass
+
+
+@migrate.command("session", cls=HelpfulCommand)
+@click.option(
+    "--source_db_url",
+    required=True,
+    help=(
+        "SQLAlchemy URL of source database in database session service, e.g."
+        " sqlite:///source.db."
+    ),
+)
+@click.option(
+    "--dest_db_url",
+    required=True,
+    help=(
+        "SQLAlchemy URL of destination database in database session service,"
+        " e.g. sqlite:///dest.db."
+    ),
+)
+@click.option(
+    "--log_level",
+    type=LOG_LEVELS,
+    default="INFO",
+    help="Optional. Set the logging level",
+)
+def cli_migrate_session(
+    *, source_db_url: str, dest_db_url: str, log_level: str
+):
+  """Migrates a session database to the latest schema version."""
+  logs.setup_adk_logger(getattr(logging, log_level.upper()))
+  try:
+    migration_runner.upgrade(source_db_url, dest_db_url)
+    click.secho("Migration check and upgrade process finished.", fg="green")
+  except Exception as e:
+    click.secho(f"Migration failed: {e}", fg="red", err=True)
+
+
 @deploy.command("agent_engine")
 @click.option(
     "--api_key",

src/google/adk/sessions/database_session_service.py

@@ -178,12 +178,9 @@ async def _prepare_tables(self):
           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
+      except Exception as e:
+        logger.error("Failed to inspect database tables: %s", e)
+        raise
 
     # Check if tables are created and create them if not
     if self._tables_created:

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

@@ -16,6 +16,7 @@
 
 import logging
 
+from sqlalchemy import create_engine as create_sync_engine
 from sqlalchemy import inspect
 from sqlalchemy import text
 
@@ -38,14 +39,16 @@ def _get_schema_version_impl(inspector, connection) -> str:
       if result:
         return result[0]
       else:
-        return LATEST_SCHEMA_VERSION
+        raise ValueError(
+            "Schema version not found in adk_internal_metadata. The database"
+            " might be malformed."
+        )
     except Exception as e:
-      logger.warning(
-          "Failed to query schema version from adk_internal_metadata,"
-          " assuming the latest schema: %s.",
+      logger.error(
+          "Failed to query schema version from adk_internal_metadata: %s.",
           e,
       )
-      return LATEST_SCHEMA_VERSION
+      raise
   # Metadata table doesn't exist, check for v0 schema.
   # V0 schema has an 'events' table with an 'actions' column.
   if inspector.has_table("events"):
@@ -57,17 +60,57 @@ def _get_schema_version_impl(inspector, connection) -> str:
             " 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."
+            " data. You can use `adk migrate session` command to migrate your"
+            " database."
         )
         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.
+      logger.error("Failed to inspect 'events' table columns: %s", e)
+      raise
+  # New database, use 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)
+
+
+def _to_sync_url(db_url: str) -> str:
+  """Removes '+driver' from SQLAlchemy URL."""
+  if "://" in db_url:
+    scheme, _, rest = db_url.partition("://")
+    if "+" in scheme:
+      dialect = scheme.split("+", 1)[0]
+      return f"{dialect}://{rest}"
+  return db_url
+
+
+def get_db_schema_version(db_url: str) -> str:
+  """Reads schema version from DB.
+
+  Checks metadata table first and then falls back to table structure.
+
+  Args:
+    db_url: The database URL.
+
+  Returns:
+    The detected schema version as a string. Returns `LATEST_SCHEMA_VERSION`
+    if it's a new database.
+  """
+  engine = None
+  try:
+    engine = create_sync_engine(_to_sync_url(db_url))
+    with engine.connect() as connection:
+      inspector = inspect(connection)
+      return _get_schema_version_impl(inspector, connection)
+  except Exception:
+    logger.warning(
+        "Failed to get schema version from database %s.",
+        db_url,
+    )
+    raise
+  finally:
+    if engine:
+      engine.dispose()