From 78d1847a0aee3969c0453269ef2129aa4f94f99b Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 22:47:33 +0000 Subject: [PATCH 1/2] fix(migration): restrict unpickling of v0 actions blobs (v1) Port of the upstream fix (PR #5866) to the v1 branch. The v0 session schema stored event actions as a pickled blob. Migration read that column straight out of the source database and handed the raw bytes to pickle.loads, so migrating a database whose contents the operator did not control could run code from that database inside the migrating process. Migration now loads those blobs through a pickle.Unpickler subclass whose find_class accepts only an explicit allowlist: builtin containers and primitives, datetime types, the ADK action models (EventActions, EventCompaction, AuthConfig, ToolConfirmation, UiWidget and the auth credential and scheme models) and the google.genai.types classes that compaction payloads need. Anything else raises pickle.UnpicklingError, which the existing handler turns into a warning and an empty EventActions() for that one event, leaving the rest of the migration to proceed. An operator who trusts the source database and needs the old behaviour for custom objects can opt back in: migrate() and migration_runner.upgrade() take allow_unsafe_unpickling, and both `adk migrate session` and the migration script take --allow_unsafe_unpickling / --allow-unsafe-unpickling. The same commit tightens the two JSON helpers in the migration path so that they accept only JSON objects. _safe_json_load returns None and _get_state_dict returns {} when the decoded value is something else. Behaviour changes for existing users: - A v0 database containing a pickled custom Python object in state_delta or another Any-typed field now migrates that event with empty actions and a logged warning, instead of reconstructing the object. Passing --allow_unsafe_unpickling restores the old behaviour. - An event whose content, metadata or transcription column holds valid JSON that is not an object used to fail model validation, which dropped the whole event with a warning. That event now migrates with the field left unset. - A state column holding valid JSON that is not an object used to be stored as-is; it is now stored as an empty dict with a warning. --- src/google/adk/cli/cli_tools_click.py | 22 +- .../migrate_from_sqlalchemy_pickle.py | 151 ++++++++- .../sessions/migration/migration_runner.py | 18 +- .../cli/utils/test_cli_tools_click.py | 47 +++ .../sessions/migration/test_migration.py | 310 ++++++++++++++++++ 5 files changed, 533 insertions(+), 15 deletions(-) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index fda251da10a..a4d2fd61a29 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/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) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index a6d1ad2a788..65a78c94012 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/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,7 +139,9 @@ 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") @@ -59,7 +149,9 @@ def _row_to_event(row: dict) -> Event: 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,8 +167,7 @@ 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) @@ -84,8 +175,17 @@ def _safe_json_load(val): 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,23 +247,31 @@ 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 @@ -172,6 +280,11 @@ def migrate(source_db_url: str, dest_db_url: str): dest_sync_url = _schema_check_utils.to_sync_url(dest_db_url) logger.info(f"Connecting to source database: {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) @@ -265,7 +378,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 +425,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) diff --git a/src/google/adk/sessions/migration/migration_runner.py b/src/google/adk/sessions/migration/migration_runner.py index c46bab2179f..1290ee67fcc 100644 --- a/src/google/adk/sessions/migration/migration_runner.py +++ b/src/google/adk/sessions/migration/migration_runner.py @@ -42,7 +42,11 @@ LATEST_VERSION = _schema_check_utils.LATEST_SCHEMA_VERSION -def upgrade(source_db_url: str, dest_db_url: str): +def upgrade( + source_db_url: str, + dest_db_url: str, + allow_unsafe_unpickling: bool = False, +) -> None: """Migrates a database from its current version to the latest version. If the source database schema is older than the latest version, this @@ -61,6 +65,9 @@ def upgrade(source_db_url: str, dest_db_url: str): source_db_url: The SQLAlchemy URL of the database to migrate from. dest_db_url: The SQLAlchemy URL of the database to migrate to. This must be different from source_db_url. + allow_unsafe_unpickling: If true, use Python's unsafe pickle loader for the + legacy pickle migration step. Only use this with a trusted source + database. Raises: RuntimeError: If source_db_url and dest_db_url are the same, or if no @@ -113,7 +120,14 @@ def upgrade(source_db_url: str, dest_db_url: str): logger.info( f"Migrating from {in_url} to {out_url} (schema v{end_version})..." ) - migrate_func(in_url, out_url) + if migrate_func is migrate_from_sqlalchemy_pickle.migrate: + migrate_func( + in_url, + out_url, + allow_unsafe_unpickling=allow_unsafe_unpickling, + ) + else: + migrate_func(in_url, out_url) logger.info("Finished migration step to schema %s.", end_version) # The output of this step becomes the input for the next step. in_url = out_url diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index ad47c9ecbe6..4e8f23c5c05 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -596,6 +596,53 @@ def test_cli_web_passes_service_uris( assert called_kwargs.get("memory_service_uri") == "rag://mycorpus" +@pytest.mark.parametrize( + "flag", + ["--allow-unsafe-unpickling", "--allow_unsafe_unpickling"], +) +def test_cli_migrate_session_allows_unsafe_unpickling_flag( + monkeypatch: pytest.MonkeyPatch, flag: str +) -> None: + calls: list[dict[str, Any]] = [] + + def fake_upgrade( + source_db_url: str, + dest_db_url: str, + *, + allow_unsafe_unpickling: bool = False, + ) -> None: + calls.append({ + "source_db_url": source_db_url, + "dest_db_url": dest_db_url, + "allow_unsafe_unpickling": allow_unsafe_unpickling, + }) + + monkeypatch.setattr( + "google.adk.sessions.migration.migration_runner.upgrade", + fake_upgrade, + ) + + result = CliRunner().invoke( + cli_tools_click.main, + [ + "migrate", + "session", + "--source_db_url", + "sqlite:///source.db", + "--dest_db_url", + "sqlite:///dest.db", + flag, + ], + ) + + assert result.exit_code == 0, (result.output, repr(result.exception)) + assert calls == [{ + "source_db_url": "sqlite:///source.db", + "dest_db_url": "sqlite:///dest.db", + "allow_unsafe_unpickling": True, + }] + + def test_cli_eval_with_eval_set_file_path( mock_load_eval_set_from_file, mock_get_root_agent, diff --git a/tests/unittests/sessions/migration/test_migration.py b/tests/unittests/sessions/migration/test_migration.py index 5d9c207f209..9f54a2e0ace 100644 --- a/tests/unittests/sessions/migration/test_migration.py +++ b/tests/unittests/sessions/migration/test_migration.py @@ -17,14 +17,23 @@ from datetime import datetime from datetime import timezone +import os +import pickle +from fastapi.openapi.models import HTTPBearer +from google.adk.auth.auth_tool import AuthConfig from google.adk.events.event_actions import EventActions +from google.adk.events.event_actions import EventCompaction +from google.adk.events.ui_widget import UiWidget from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle as mfsp from google.adk.sessions.schemas import v0 from google.adk.sessions.schemas import v1 +from google.adk.tools.tool_confirmation import ToolConfirmation +from google.genai import types import pytest from sqlalchemy import create_engine +from sqlalchemy import text from sqlalchemy.orm import sessionmaker @@ -184,6 +193,307 @@ def test_migrate_from_sqlalchemy_pickle(tmp_path): dest_session.close() +def test_migrate_from_sqlalchemy_pickle_preserves_safe_actions_pickle(tmp_path): + """Migration should preserve normal v0 EventActions pickle payloads.""" + source_db_path = tmp_path / "source_pickle_safe_actions.db" + dest_db_path = tmp_path / "dest_json_safe_actions.db" + source_db_url = f"sqlite:///{source_db_path}" + dest_db_url = f"sqlite:///{dest_db_path}" + + source_engine = create_engine(source_db_url) + v0.Base.metadata.create_all(source_engine) + SourceSession = sessionmaker(bind=source_engine) + + now = datetime.now(timezone.utc) + with SourceSession() as source_session: + source_session.add( + v0.StorageSession( + app_name="app1", + user_id="user1", + id="session1", + state={}, + create_time=now, + update_time=now, + ) + ) + source_session.commit() + + actions = EventActions( + state_delta={"skey": "updated"}, + artifact_delta={"artifact.txt": 2}, + ) + source_session.add( + v0.StorageEvent( + id="event1", + app_name="app1", + user_id="user1", + session_id="session1", + invocation_id="invoke1", + author="user", + actions=actions, + timestamp=now, + ) + ) + source_session.commit() + + mfsp.migrate(source_db_url, dest_db_url) + + dest_engine = create_engine(dest_db_url) + DestSession = sessionmaker(bind=dest_engine) + with DestSession() as dest_session: + event_res = dest_session.query(v1.StorageEvent).first() + assert event_res is not None + assert event_res.event_data["actions"]["state_delta"] == {"skey": "updated"} + assert event_res.event_data["actions"]["artifact_delta"] == { + "artifact.txt": 2 + } + + +def test_migrate_from_sqlalchemy_pickle_preserves_nested_safe_actions_pickle( + tmp_path, +): + """Migration should allow standard nested EventActions models.""" + source_db_path = tmp_path / "source_pickle_nested_actions.db" + dest_db_path = tmp_path / "dest_json_nested_actions.db" + source_db_url = f"sqlite:///{source_db_path}" + dest_db_url = f"sqlite:///{dest_db_path}" + + source_engine = create_engine(source_db_url) + v0.Base.metadata.create_all(source_engine) + SourceSession = sessionmaker(bind=source_engine) + + now = datetime.now(timezone.utc) + with SourceSession() as source_session: + source_session.add( + v0.StorageSession( + app_name="app1", + user_id="user1", + id="session1", + state={}, + create_time=now, + update_time=now, + ) + ) + source_session.commit() + + actions = EventActions( + requested_auth_configs={ + "fc-auth": AuthConfig(auth_scheme=HTTPBearer()) + }, + requested_tool_confirmations={ + "fc-confirm": ToolConfirmation(hint="Authorize execution?") + }, + compaction=EventCompaction( + start_timestamp=1.0, + end_timestamp=2.0, + compacted_content=types.Content( + parts=[types.Part(text="summary")], + role="model", + ), + ), + ) + source_session.add( + v0.StorageEvent( + id="event1", + app_name="app1", + user_id="user1", + session_id="session1", + invocation_id="invoke1", + author="user", + actions=actions, + timestamp=now, + ) + ) + source_session.commit() + + mfsp.migrate(source_db_url, dest_db_url) + + dest_engine = create_engine(dest_db_url) + DestSession = sessionmaker(bind=dest_engine) + with DestSession() as dest_session: + event_res = dest_session.query(v1.StorageEvent).first() + assert event_res is not None + actions_data = event_res.event_data["actions"] + assert "fc-auth" in actions_data["requested_auth_configs"] + assert ( + actions_data["requested_tool_confirmations"]["fc-confirm"]["hint"] + == "Authorize execution?" + ) + assert ( + actions_data["compaction"]["compacted_content"]["parts"][0]["text"] + == "summary" + ) + + +def test_restricted_actions_unpickler_allows_datetime_state_delta(): + """Standard timestamp objects in action deltas should migrate by default.""" + last_seen = datetime(2026, 1, 1, 12, 30, tzinfo=timezone.utc) + actions = EventActions(state_delta={"last_seen": last_seen}) + + loaded_actions = mfsp._restricted_pickle_loads(pickle.dumps(actions)) + + assert isinstance(loaded_actions, EventActions) + assert loaded_actions.state_delta["last_seen"] == last_seen + + +def test_restricted_actions_unpickler_allows_ui_widgets(): + """Standard UI widget action metadata should migrate by default.""" + actions = EventActions( + render_ui_widgets=[ + UiWidget( + id="widget-1", + provider="mcp", + payload={"resource_uri": "ui://widget"}, + ) + ] + ) + + loaded_actions = mfsp._restricted_pickle_loads(pickle.dumps(actions)) + + assert isinstance(loaded_actions, EventActions) + assert loaded_actions.render_ui_widgets == actions.render_ui_widgets + + +def test_migrate_from_sqlalchemy_pickle_ignores_non_object_json_fields(): + """Event JSON model fields should only decode object payloads.""" + event = mfsp._row_to_event({ + "id": "event-list-content", + "invocation_id": "invoke1", + "author": "user", + "timestamp": datetime(2026, 1, 1, tzinfo=timezone.utc), + "content": "[1, 2, 3]", + }) + + assert event.content is None + + +def test_migrate_from_sqlalchemy_pickle_blocks_unsafe_actions_pickle( + tmp_path, monkeypatch +): + """Migration should not execute arbitrary globals from a pickled actions blob.""" + monkeypatch.delenv("ADK_MIGRATION_PICKLE_RCE", raising=False) + + source_db_path = tmp_path / "source_pickle_unsafe_actions.db" + dest_db_path = tmp_path / "dest_json_unsafe_actions.db" + source_db_url = f"sqlite:///{source_db_path}" + dest_db_url = f"sqlite:///{dest_db_path}" + + source_engine = create_engine(source_db_url) + v0.Base.metadata.create_all(source_engine) + SourceSession = sessionmaker(bind=source_engine) + + # Populate source DB with a valid session row to satisfy the FK constraint, + # then insert a malicious pickled actions blob directly as raw bytes. + now = datetime.now(timezone.utc) + with SourceSession() as source_session: + source_session.add( + v0.StorageSession( + app_name="app1", + user_id="user1", + id="session1", + state={}, + create_time=now, + update_time=now, + ) + ) + source_session.commit() + + class Evil: + + def __reduce__(self): + # This is intentionally non-destructive: it only sets an env var. + return ( + exec, + ("import os; os.environ['ADK_MIGRATION_PICKLE_RCE']='1'",), + ) + + source_session.execute( + text( + "INSERT INTO events (id, app_name, user_id, session_id," + " invocation_id, author, actions, timestamp) VALUES (:id," + " :app_name, :user_id, :session_id, :invocation_id, :author," + " :actions, :timestamp)" + ), + { + "id": "event1", + "app_name": "app1", + "user_id": "user1", + "session_id": "session1", + "invocation_id": "invoke1", + "author": "user", + "actions": pickle.dumps(Evil()), + "timestamp": now, + }, + ) + source_session.commit() + + mfsp.migrate(source_db_url, dest_db_url) + + assert os.environ.get("ADK_MIGRATION_PICKLE_RCE") is None + + +def test_migrate_from_sqlalchemy_pickle_allows_unsafe_actions_pickle_when_opted_in( + tmp_path, monkeypatch +): + """Unsafe pickle loading should require an explicit migration opt-in.""" + monkeypatch.delenv("ADK_MIGRATION_PICKLE_RCE", raising=False) + + source_db_path = tmp_path / "source_pickle_unsafe_opt_in_actions.db" + dest_db_path = tmp_path / "dest_json_unsafe_opt_in_actions.db" + source_db_url = f"sqlite:///{source_db_path}" + dest_db_url = f"sqlite:///{dest_db_path}" + + source_engine = create_engine(source_db_url) + v0.Base.metadata.create_all(source_engine) + SourceSession = sessionmaker(bind=source_engine) + + now = datetime.now(timezone.utc) + with SourceSession() as source_session: + source_session.add( + v0.StorageSession( + app_name="app1", + user_id="user1", + id="session1", + state={}, + create_time=now, + update_time=now, + ) + ) + source_session.commit() + + class Evil: + + def __reduce__(self): + return ( + exec, + ("import os; os.environ['ADK_MIGRATION_PICKLE_RCE']='1'",), + ) + + source_session.execute( + text( + "INSERT INTO events (id, app_name, user_id, session_id," + " invocation_id, author, actions, timestamp) VALUES (:id," + " :app_name, :user_id, :session_id, :invocation_id, :author," + " :actions, :timestamp)" + ), + { + "id": "event1", + "app_name": "app1", + "user_id": "user1", + "session_id": "session1", + "invocation_id": "invoke1", + "author": "user", + "actions": pickle.dumps(Evil()), + "timestamp": now, + }, + ) + source_session.commit() + + mfsp.migrate(source_db_url, dest_db_url, allow_unsafe_unpickling=True) + + assert os.environ.get("ADK_MIGRATION_PICKLE_RCE") == "1" + + def test_migrate_from_sqlalchemy_pickle_with_async_driver_urls(tmp_path): """Tests that migration works with async driver URLs (fixes issue #4176). From c0332d9526e627d1af9e6e394d3e0b3d9abc55c2 Mon Sep 17 00:00:00 2001 From: George Weale Date: Mon, 17 Aug 2026 22:50:44 +0000 Subject: [PATCH 2/2] fix: redact database password from session service errors and logs (v1) Port of the upstream fix (PR #6485) to the v1 branch. A database URL carries its password in the userinfo component, and the session code interpolated that URL as-is into places that routinely end up in application logs and tracebacks. DatabaseSessionService put it in all three of its engine-creation ValueError messages, the schema-version check logged it in a warning, and the migration entry points logged it on every run, not just on failure: both connect lines in the pickle migration, the connect line in the sqlite migration, and the already-up-to-date and per-step lines in the migration runner. _schema_check_utils now has a _redact_db_url helper, and those call sites go through it. It parses the URL with SQLAlchemy's make_url, replaces every query-parameter value with REDACTED, and renders the result with hide_password=True. Query values are masked wholesale because drivers accept secrets as query parameters under names ADK cannot enumerate. The helper runs while an error is already being reported, so it catches everything and returns the fixed string "" rather than raising or echoing a URL it could not parse. Behaviour change for existing users: these error messages and log lines no longer contain the full connection string. The password shows as ***, each query-parameter value shows as REDACTED, and a URL that make_url cannot parse is replaced by the placeholder. Anyone grepping logs for a connection string, or parsing the ValueError text, will see different output. --- .../adk/sessions/database_session_service.py | 7 +- .../sessions/migration/_schema_check_utils.py | 24 +++- .../migrate_from_sqlalchemy_pickle.py | 10 +- .../migrate_from_sqlalchemy_sqlite.py | 5 +- .../sessions/migration/migration_runner.py | 10 +- .../sessions/migration/test_migration.py | 116 ++++++++++++++++++ .../sessions/test_session_service.py | 43 +++++++ 7 files changed, 205 insertions(+), 10 deletions(-) diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index d033f1f2347..70842e0eb1f 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/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 diff --git a/src/google/adk/sessions/migration/_schema_check_utils.py b/src/google/adk/sessions/migration/_schema_check_utils.py index a6bc8a546a6..c02425ce651 100644 --- a/src/google/adk/sessions/migration/_schema_check_utils.py +++ b/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 = "" +_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: diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py index 65a78c94012..fd391e83ffd 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -279,7 +279,10 @@ def migrate( 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" @@ -292,7 +295,10 @@ def migrate( 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) diff --git a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_sqlite.py b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_sqlite.py index dbd2cef3ba4..b2a7c4240e8 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_sqlite.py +++ b/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( diff --git a/src/google/adk/sessions/migration/migration_runner.py b/src/google/adk/sessions/migration/migration_runner.py index 1290ee67fcc..d7b57d82e16 100644 --- a/src/google/adk/sessions/migration/migration_runner.py +++ b/src/google/adk/sessions/migration/migration_runner.py @@ -82,8 +82,9 @@ def upgrade( current_version = _schema_check_utils.get_db_schema_version(source_db_url) if current_version == LATEST_VERSION: logger.info( - f"Database {source_db_url} is already at latest version" - f" {LATEST_VERSION}. No migration needed." + "Database %s is already at latest version %s. No migration needed.", + _schema_check_utils._redact_db_url(source_db_url), + LATEST_VERSION, ) return @@ -118,7 +119,10 @@ def upgrade( logger.debug("Created temp db %s for step %d", out_url, i + 1) logger.info( - f"Migrating from {in_url} to {out_url} (schema v{end_version})..." + "Migrating from %s to %s (schema v%s)...", + _schema_check_utils._redact_db_url(in_url), + _schema_check_utils._redact_db_url(out_url), + end_version, ) if migrate_func is migrate_from_sqlalchemy_pickle.migrate: migrate_func( diff --git a/tests/unittests/sessions/migration/test_migration.py b/tests/unittests/sessions/migration/test_migration.py index 9f54a2e0ace..402c554a50a 100644 --- a/tests/unittests/sessions/migration/test_migration.py +++ b/tests/unittests/sessions/migration/test_migration.py @@ -17,8 +17,10 @@ from datetime import datetime from datetime import timezone +import logging import os import pickle +from unittest import mock from fastapi.openapi.models import HTTPBearer from google.adk.auth.auth_tool import AuthConfig @@ -27,6 +29,8 @@ from google.adk.events.ui_widget import UiWidget from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle as mfsp +from google.adk.sessions.migration import migrate_from_sqlalchemy_sqlite as mfss +from google.adk.sessions.migration import migration_runner from google.adk.sessions.schemas import v0 from google.adk.sessions.schemas import v1 from google.adk.tools.tool_confirmation import ToolConfirmation @@ -114,6 +118,118 @@ def test_to_sync_url_empty_string(self): assert _schema_check_utils.to_sync_url("") == "" +class TestRedactDbUrl: + """Tests for the _redact_db_url function.""" + + def test_password_is_masked(self): + redacted = _schema_check_utils._redact_db_url( + "postgresql+asyncpg://user:sup3r-s3cret@host:5432/db" + ) + assert redacted == "postgresql+asyncpg://user:***@host:5432/db" + + def test_unparseable_url_falls_back_to_placeholder(self): + """Redaction runs while reporting an error, so it must never raise.""" + assert ( + _schema_check_utils._redact_db_url("definitely not a url sup3r-s3cret") + == "" + ) + + def test_query_parameter_values_are_masked(self): + """Drivers accept secrets as query parameters, so every value is masked.""" + redacted = _schema_check_utils._redact_db_url( + "postgresql://user@host:5432/db?password=sup3r-s3cret&sslmode=require" + ) + assert redacted == ( + "postgresql://user@host:5432/db?password=REDACTED&sslmode=REDACTED" + ) + + def test_schema_version_failure_warning_hides_password(self, caplog): + db_url = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/db" + + with mock.patch.object( + _schema_check_utils, + "create_sync_engine", + side_effect=RuntimeError("boom"), + ): + with caplog.at_level(logging.WARNING): + with pytest.raises(RuntimeError): + _schema_check_utils.get_db_schema_version(db_url) + + assert "sup3r-s3cret" not in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/db" in caplog.text + + +_SOURCE_URL = "postgresql+asyncpg://user:sup3r-s3cret@host:5432/src" +_DEST_URL = "postgresql+asyncpg://user:0ther-s3cret@host:5432/dst" + + +class TestMigrationLogsHidePassword: + """These entry points log their URLs on every run, not only on failure.""" + + def test_pickle_migration_connect_logs_are_redacted(self, caplog): + with mock.patch.object( + mfsp, + "create_engine", + side_effect=[mock.MagicMock(), RuntimeError("boom")], + ): + with caplog.at_level(logging.INFO): + with pytest.raises(RuntimeError): + mfsp.migrate(_SOURCE_URL, _DEST_URL) + + assert "sup3r-s3cret" not in caplog.text + assert "0ther-s3cret" not in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text + + def test_sqlite_migration_connect_log_is_redacted(self, caplog, tmp_path): + with mock.patch.object( + mfss, "create_engine", side_effect=RuntimeError("boom") + ): + with caplog.at_level(logging.INFO): + with pytest.raises(SystemExit): + mfss.migrate(_SOURCE_URL, str(tmp_path / "dest.db")) + + assert "sup3r-s3cret" not in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text + + def test_runner_up_to_date_log_is_redacted(self, caplog): + with mock.patch.object( + _schema_check_utils, + "get_db_schema_version", + return_value=migration_runner.LATEST_VERSION, + ): + with caplog.at_level(logging.INFO): + migration_runner.upgrade(_SOURCE_URL, _DEST_URL) + + assert "sup3r-s3cret" not in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text + + def test_runner_migration_step_log_is_redacted(self, caplog): + mock_migrate = mock.Mock() + with mock.patch.object( + _schema_check_utils, + "get_db_schema_version", + return_value=_schema_check_utils.SCHEMA_VERSION_0_PICKLE, + ): + with mock.patch.dict( + migration_runner.MIGRATIONS, + { + _schema_check_utils.SCHEMA_VERSION_0_PICKLE: ( + _schema_check_utils.SCHEMA_VERSION_1_JSON, + mock_migrate, + ) + }, + ): + with caplog.at_level(logging.INFO): + migration_runner.upgrade(_SOURCE_URL, _DEST_URL) + + mock_migrate.assert_called_once_with(_SOURCE_URL, _DEST_URL) + assert "sup3r-s3cret" not in caplog.text + assert "0ther-s3cret" not in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/src" in caplog.text + assert "postgresql+asyncpg://user:***@host:5432/dst" in caplog.text + + def test_migrate_from_sqlalchemy_pickle(tmp_path): """Tests for migrate_from_sqlalchemy_pickle.""" source_db_path = tmp_path / "source_pickle.db" diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index 02f5159a453..559b0ace68f 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -33,6 +33,7 @@ from google.genai import types import pytest from sqlalchemy import delete +from sqlalchemy.exc import ArgumentError class SessionServiceType(enum.Enum): @@ -1650,3 +1651,45 @@ async def tracking_fn(**kwargs): finally: database_session_service._select_required_state = original_fn await service.close() + + +@pytest.mark.parametrize( + 'raised_error', + [ + RuntimeError('boom'), + ArgumentError('bad argument'), + ImportError('no driver'), + ], +) +def test_database_session_service_engine_error_hides_password(raised_error): + """Engine creation errors must not put the DB password in the message.""" + password = 'sup3r-s3cret' + db_url = f'postgresql+asyncpg://user:{password}@localhost:5432/db' + + with mock.patch.object( + database_session_service, + 'create_async_engine', + side_effect=raised_error, + ): + with pytest.raises(ValueError) as exc_info: + DatabaseSessionService(db_url) + + message = str(exc_info.value) + assert password not in message + # The redacted URL is still there, so the error stays diagnosable. + assert 'postgresql+asyncpg://user:***@localhost:5432/db' in message + + +def test_database_session_service_malformed_url_reports_usable_error(): + """A URL too malformed to parse still yields a usable, leak-free error.""" + # make_url() itself rejects this, so redaction cannot parse it either and + # must fall back to a placeholder rather than echoing the raw string. + db_url = 'definitely not a url sup3r-s3cret' + + with pytest.raises(ValueError) as exc_info: + DatabaseSessionService(db_url) + + message = str(exc_info.value) + assert 'sup3r-s3cret' not in message + assert 'Invalid database URL format or argument' in message + assert isinstance(exc_info.value.__cause__, ArgumentError)