From 6575f7847118731cb5a4c683e768da08561e9d52 Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Sun, 23 Aug 2026 10:55:30 -0500 Subject: [PATCH] fix(sessions): Port PostgreSQL binary event actions migration fix to v1 --- .../migrate_from_sqlalchemy_pickle.py | 9 +++++-- .../sessions/migration/test_migration.py | 26 +++++++++++++++++++ 2 files changed, 33 insertions(+), 2 deletions(-) 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 fd391e83ff..90ac5582ba 100644 --- a/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py +++ b/src/google/adk/sessions/migration/migrate_from_sqlalchemy_pickle.py @@ -148,9 +148,14 @@ def _row_to_event( actions = None if actions_val is not None: try: - if isinstance(actions_val, bytes): + # The source rows are read with raw SQL, so SQLAlchemy has no column + # type to coerce with and whatever the driver produced for the binary + # column arrives here untouched. psycopg2 produces a memoryview rather + # than bytes, so match every bytes-like form instead of one driver's. + if isinstance(actions_val, (bytes, bytearray, memoryview)): actions = _restricted_pickle_loads( - actions_val, allow_unsafe_unpickling=allow_unsafe_unpickling + bytes(actions_val), + allow_unsafe_unpickling=allow_unsafe_unpickling, ) else: # for spanner - it might return object directly actions = actions_val diff --git a/tests/unittests/sessions/migration/test_migration.py b/tests/unittests/sessions/migration/test_migration.py index 402c554a50..2fa65282bd 100644 --- a/tests/unittests/sessions/migration/test_migration.py +++ b/tests/unittests/sessions/migration/test_migration.py @@ -483,6 +483,32 @@ def test_migrate_from_sqlalchemy_pickle_ignores_non_object_json_fields(): assert event.content is None +@pytest.mark.parametrize("as_binary", [bytes, bytearray, memoryview]) +def test_migrate_from_sqlalchemy_pickle_reads_every_binary_column_type( + as_binary, +): + """Pickled actions must survive whichever binary type the driver returns. + + Events are read with raw SQL, so SQLAlchemy has no column type to coerce + with and the driver's own representation reaches the migration: psycopg2 + returns a memoryview rather than bytes. Treating that as "some other + backend handed us an object" replaced the actions with an empty one while + the migration still reported success. + """ + actions = EventActions(state_delta={"skey": 4}, escalate=True) + + event = mfsp._row_to_event({ + "id": "event-binary-actions", + "invocation_id": "invoke1", + "author": "user", + "timestamp": datetime(2026, 1, 1, tzinfo=timezone.utc), + "actions": as_binary(pickle.dumps(actions)), + }) + + assert event.actions.state_delta == {"skey": 4} + assert event.actions.escalate is True + + def test_migrate_from_sqlalchemy_pickle_blocks_unsafe_actions_pickle( tmp_path, monkeypatch ):