fix: Port migration unpickle allowlist and database URL redaction to v1

src/google/adk/cli/cli_tools_click.py

@@ -2005,15 +2005,33 @@ def migrate():
     default="INFO",
     help="Optional. Set the logging level",
 )
+@click.option(  # type: ignore[untyped-decorator]
+    "--allow-unsafe-unpickling",
+    "--allow_unsafe_unpickling",
+    is_flag=True,
+    default=False,
+    help=(
+        "Optional. Allow unsafe pickle loading for trusted legacy session"
+        " databases."
+    ),
+)
 def cli_migrate_session(
-    *, source_db_url: str, dest_db_url: str, log_level: str
+    *,
+    source_db_url: str,
+    dest_db_url: str,
+    log_level: str,
+    allow_unsafe_unpickling: bool,
 ):
   """Migrates a session database to the latest schema version."""
   logs.setup_adk_logger(getattr(logging, log_level.upper()))
   try:
     from ..sessions.migration import migration_runner
 
-    migration_runner.upgrade(source_db_url, dest_db_url)
+    migration_runner.upgrade(
+        source_db_url,
+        dest_db_url,
+        allow_unsafe_unpickling=allow_unsafe_unpickling,
+    )
     click.secho("Migration check and upgrade process finished.", fg="green")
   except Exception as e:
     click.secho(f"Migration failed: {e}", fg="red", err=True)

src/google/adk/sessions/database_session_service.py

@@ -213,16 +213,17 @@ def __init__(self, db_url: str, **kwargs: Any):
         event.listen(db_engine.sync_engine, "connect", _set_sqlite_pragma)
 
     except Exception as e:
+      redacted_url = _schema_check_utils._redact_db_url(db_url)
       if isinstance(e, ArgumentError):
         raise ValueError(
-            f"Invalid database URL format or argument '{db_url}'."
+            f"Invalid database URL format or argument '{redacted_url}'."
         ) from e
       if isinstance(e, ImportError):
         raise ValueError(
-            f"Database related module not found for URL '{db_url}'."
+            f"Database related module not found for URL '{redacted_url}'."
         ) from e
       raise ValueError(
-          f"Failed to create database engine for URL '{db_url}'"
+          f"Failed to create database engine for URL '{redacted_url}'"
       ) from e
 
     self.db_engine: AsyncEngine = db_engine

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

@@ -20,9 +20,13 @@
 from sqlalchemy import create_engine as create_sync_engine
 from sqlalchemy import inspect
 from sqlalchemy import text
+from sqlalchemy.engine import make_url
 
 logger = logging.getLogger("google_adk." + __name__)
 
+_UNPARSEABLE_DB_URL = "<unparseable database URL>"
+_REDACTED_QUERY_VALUE = "REDACTED"
+
 SCHEMA_VERSION_KEY = "schema_version"
 SCHEMA_VERSION_0_PICKLE = "0"
 SCHEMA_VERSION_1_JSON = "1"
@@ -112,6 +116,24 @@ def to_sync_url(db_url: str) -> str:
   return db_url
 
 
+def _redact_db_url(db_url: str) -> str:
+  """Returns the URL with its credentials masked, for logs and error messages.
+
+  A database URL carries the password in the userinfo component, and drivers
+  also accept secrets as query parameters, so every query value is masked
+  rather than only the ones with a recognizable name. Redaction happens while
+  an error is being reported, so it never raises: an unparseable URL yields a
+  fixed placeholder rather than the original string.
+  """
+  try:
+    url = make_url(db_url)
+    if url.query:
+      url = url.set(query={key: _REDACTED_QUERY_VALUE for key in url.query})
+    return str(url.render_as_string(hide_password=True))
+  except Exception:  # pylint: disable=broad-except
+    return _UNPARSEABLE_DB_URL
+
+
 def get_db_schema_version(db_url: str) -> str:
   """Reads schema version from DB.
 
@@ -133,7 +155,7 @@ def get_db_schema_version(db_url: str) -> str:
   except Exception:
     logger.warning(
         "Failed to get schema version from database %s.",
-        db_url,
+        _redact_db_url(db_url),
     )
     raise
   finally:

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

@@ -18,6 +18,7 @@
 import argparse
 from datetime import datetime
 from datetime import timezone
+import io
 import json
 import logging
 import pickle
@@ -37,6 +38,93 @@
 
 logger = logging.getLogger("google_adk." + __name__)
 
+_ALLOWED_PICKLE_GLOBALS: set[tuple[str, str]] = {
+    # Builtin containers/primitives.
+    ("builtins", "dict"),
+    ("builtins", "list"),
+    ("builtins", "set"),
+    ("builtins", "tuple"),
+    ("builtins", "str"),
+    ("builtins", "bytes"),
+    ("builtins", "bytearray"),
+    ("builtins", "int"),
+    ("builtins", "float"),
+    ("builtins", "bool"),
+    ("datetime", "datetime"),
+    ("datetime", "timedelta"),
+    ("datetime", "timezone"),
+    # Expected pickled payload for v0 session schema events.
+    ("fastapi.openapi.models", "APIKey"),
+    ("fastapi.openapi.models", "APIKeyIn"),
+    ("fastapi.openapi.models", "HTTPBase"),
+    ("fastapi.openapi.models", "HTTPBearer"),
+    ("fastapi.openapi.models", "OAuth2"),
+    ("fastapi.openapi.models", "OAuthFlow"),
+    ("fastapi.openapi.models", "OAuthFlowAuthorizationCode"),
+    ("fastapi.openapi.models", "OAuthFlowClientCredentials"),
+    ("fastapi.openapi.models", "OAuthFlowImplicit"),
+    ("fastapi.openapi.models", "OAuthFlowPassword"),
+    ("fastapi.openapi.models", "OAuthFlows"),
+    ("fastapi.openapi.models", "OpenIdConnect"),
+    ("fastapi.openapi.models", "SecurityBase"),
+    ("fastapi.openapi.models", "SecurityScheme"),
+    ("fastapi.openapi.models", "SecuritySchemeType"),
+    ("google.adk.auth.auth_credential", "AuthCredential"),
+    ("google.adk.auth.auth_credential", "AuthCredentialTypes"),
+    ("google.adk.auth.auth_credential", "HttpAuth"),
+    ("google.adk.auth.auth_credential", "HttpCredentials"),
+    ("google.adk.auth.auth_credential", "OAuth2Auth"),
+    ("google.adk.auth.auth_credential", "ServiceAccountCredential"),
+    ("google.adk.auth.auth_schemes", "CustomAuthScheme"),
+    ("google.adk.auth.auth_schemes", "ExtendedOAuth2"),
+    ("google.adk.auth.auth_schemes", "OAuthGrantType"),
+    ("google.adk.auth.auth_schemes", "OpenIdConnectWithConfig"),
+    ("google.adk.auth.auth_tool", "AuthConfig"),
+    ("google.adk.events.event_actions", "EventActions"),
+    ("google.adk.events.event_actions", "EventCompaction"),
+    ("google.adk.events.ui_widget", "UiWidget"),
+    ("google.adk.tools.tool_confirmation", "ToolConfirmation"),
+    ("google.genai.types", "Blob"),
+    ("google.genai.types", "CodeExecutionResult"),
+    ("google.genai.types", "Content"),
+    ("google.genai.types", "ExecutableCode"),
+    ("google.genai.types", "FileData"),
+    ("google.genai.types", "FunctionCall"),
+    ("google.genai.types", "FunctionResponse"),
+    ("google.genai.types", "FunctionResponseBlob"),
+    ("google.genai.types", "FunctionResponseFileData"),
+    ("google.genai.types", "FunctionResponsePart"),
+    ("google.genai.types", "Part"),
+    ("google.genai.types", "PartMediaResolution"),
+    ("google.genai.types", "VideoMetadata"),
+}
+
+
+class _RestrictedUnpickler(pickle.Unpickler):
+  """Restricted unpickler for migrating legacy v0 schema actions.
+
+  The v0 session schema stored `EventActions` as a pickled blob. During
+  migration we treat the raw bytes read from the source DB as untrusted input
+  and only allow the minimum set of safe globals needed to reconstruct
+  `EventActions`.
+  """
+
+  def find_class(self, module: str, name: str) -> Any:  # noqa: ANN001
+    if (module, name) in _ALLOWED_PICKLE_GLOBALS:
+      return super().find_class(module, name)
+    raise pickle.UnpicklingError(
+        f"Blocked global during migration unpickle: {module}.{name}"
+    )
+
+
+def _restricted_pickle_loads(
+    data: bytes, *, allow_unsafe_unpickling: bool = False
+) -> Any:
+  """Load a pickle payload using the restricted unpickler by default."""
+  if allow_unsafe_unpickling:
+    return pickle.loads(data)
+  return _RestrictedUnpickler(io.BytesIO(data)).load()
+
 
 def _to_datetime_obj(val: Any) -> datetime | Any:
   """Converts string to datetime if needed."""
@@ -51,15 +139,19 @@ def _to_datetime_obj(val: Any) -> datetime | Any:
   return val
 
 
-def _row_to_event(row: dict) -> Event:
+def _row_to_event(
+    row: dict[str, Any], *, allow_unsafe_unpickling: bool = False
+) -> Event:
   """Converts event row (dict) to event object, handling missing columns and deserializing."""
 
   actions_val = row.get("actions")
   actions = None
   if actions_val is not None:
     try:
       if isinstance(actions_val, bytes):
-        actions = pickle.loads(actions_val)
+        actions = _restricted_pickle_loads(
+            actions_val, allow_unsafe_unpickling=allow_unsafe_unpickling
+        )
       else:  # for spanner - it might return object directly
         actions = actions_val
     except Exception as e:
@@ -75,17 +167,25 @@ def _row_to_event(row: dict) -> Event:
   else:
     actions = EventActions()
 
-  def _safe_json_load(val):
-    data = None
+  def _safe_json_load(val: Any) -> dict[str, Any] | None:
     if isinstance(val, str):
       try:
         data = json.loads(val)
       except json.JSONDecodeError:
         logger.warning(f"Failed to decode JSON for event {row.get('id')}")
         return None
     elif isinstance(val, dict):
-      data = val  # for postgres JSONB
-    return data
+      return val  # for postgres JSONB
+    else:
+      return None
+
+    if isinstance(data, dict):
+      return data
+    logger.warning(
+        f"Expected JSON object for event {row.get('id')}, got"
+        f" {type(data).__name__}."
+    )
+    return None
 
   content_dict = _safe_json_load(row.get("content"))
   grounding_metadata_dict = _safe_json_load(row.get("grounding_metadata"))
@@ -147,39 +247,58 @@ def _safe_json_load(val):
   )
 
 
-def _get_state_dict(state_val: Any) -> dict:
+def _get_state_dict(state_val: Any) -> dict[str, Any]:
   """Safely load dict from JSON string or return dict if already dict."""
   if isinstance(state_val, dict):
     return state_val
   if isinstance(state_val, str):
     try:
-      return json.loads(state_val)
+      data = json.loads(state_val)
     except json.JSONDecodeError:
       logger.warning(
           "Failed to parse state JSON string, defaulting to empty dict."
       )
       return {}
+    if isinstance(data, dict):
+      return data
+    logger.warning("State JSON was not an object, defaulting to empty dict.")
+    return {}
   return {}
 
 
 # --- Migration Logic ---
-def migrate(source_db_url: str, dest_db_url: str):
+def migrate(
+    source_db_url: str,
+    dest_db_url: str,
+    allow_unsafe_unpickling: bool = False,
+) -> None:
   """Migrates data from old pickle schema to new JSON schema."""
   # Convert async driver URLs to sync URLs for SQLAlchemy's synchronous engine.
   # This allows users to provide URLs like 'postgresql+asyncpg://...' and have
   # them automatically converted to 'postgresql://...' for migration.
   source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
   dest_sync_url = _schema_check_utils.to_sync_url(dest_db_url)
 
-  logger.info(f"Connecting to source database: {source_db_url}")
+  logger.info(
+      "Connecting to source database: %s",
+      _schema_check_utils._redact_db_url(source_db_url),
+  )
+  if allow_unsafe_unpickling:
+    logger.warning(
+        "Unsafe pickle migration mode is enabled. Only use this with a trusted"
+        " source database."
+    )
   try:
     source_engine = create_engine(source_sync_url)
     SourceSession = sessionmaker(bind=source_engine)
   except Exception as e:
     logger.error(f"Failed to connect to source database: {e}")
     raise RuntimeError(f"Failed to connect to source database: {e}") from e
 
-  logger.info(f"Connecting to destination database: {dest_db_url}")
+  logger.info(
+      "Connecting to destination database: %s",
+      _schema_check_utils._redact_db_url(dest_db_url),
+  )
   try:
     dest_engine = create_engine(dest_sync_url)
     v1.Base.metadata.create_all(dest_engine)
@@ -265,7 +384,10 @@ def migrate(source_db_url: str, dest_db_url: str):
             text("SELECT * FROM events")
         ).mappings():
           try:
-            event_obj = _row_to_event(dict(row))
+            event_obj = _row_to_event(
+                dict(row),
+                allow_unsafe_unpickling=allow_unsafe_unpickling,
+            )
             new_event = v1.StorageEvent(
                 id=event_obj.id,
                 app_name=row["app_name"],
@@ -309,9 +431,22 @@ def migrate(source_db_url: str, dest_db_url: str):
       required=True,
       help="SQLAlchemy URL of destination database",
   )
+  parser.add_argument(
+      "--allow_unsafe_unpickling",
+      "--allow-unsafe-unpickling",
+      action="store_true",
+      help=(
+          "Allow legacy pickle payloads to use Python's unsafe pickle loader."
+          " Only use this with a trusted source database."
+      ),
+  )
   args = parser.parse_args()
   try:
-    migrate(args.source_db_url, args.dest_db_url)
+    migrate(
+        args.source_db_url,
+        args.dest_db_url,
+        allow_unsafe_unpickling=args.allow_unsafe_unpickling,
+    )
   except Exception as e:
     logger.error(f"Migration failed: {e}")
     sys.exit(1)

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

@@ -38,7 +38,10 @@ def migrate(source_db_url: str, dest_db_path: str):
   # them automatically converted to 'sqlite://...' for migration.
   source_sync_url = _schema_check_utils.to_sync_url(source_db_url)
 
-  logger.info(f"Connecting to source database: {source_db_url}")
+  logger.info(
+      "Connecting to source database: %s",
+      _schema_check_utils._redact_db_url(source_db_url),
+  )
   try:
     engine = create_engine(source_sync_url)
     v0_schema.Base.metadata.create_all(