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/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 a6d1ad2a788..fd391e83ffd 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 @@ -171,7 +279,15 @@ def migrate(source_db_url: str, dest_db_url: str): 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) @@ -179,7 +295,10 @@ def migrate(source_db_url: str, dest_db_url: str): 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) 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 c46bab2179f..d7b57d82e16 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 @@ -75,8 +82,9 @@ def upgrade(source_db_url: str, dest_db_url: str): 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 @@ -111,9 +119,19 @@ def upgrade(source_db_url: str, dest_db_url: str): 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, ) - 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..402c554a50a 100644 --- a/tests/unittests/sessions/migration/test_migration.py +++ b/tests/unittests/sessions/migration/test_migration.py @@ -17,14 +17,27 @@ 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 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.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 +from google.genai import types import pytest from sqlalchemy import create_engine +from sqlalchemy import text from sqlalchemy.orm import sessionmaker @@ -105,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" @@ -184,6 +309,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). 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)