fix: Update DynamicPickleType to support MySQL dialect

Merge https://github.com/google/adk-python/pull/3282

The `process_bind_param` and `process_result_value` methods in the `DynamicPickleType` class have been modified to handle MySQL dialect in addition to Spanner. This change ensures that pickled values are correctly processed for both database types.

**Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.**

### Link to Issue or Description of Change

**1. Link to an existing issue (if applicable):**

- Closes: #3283

**2. Or, if no issue exists, describe the change:**

_If applicable, please follow the issue templates to provide as much detail as
possible._

**Problem:**
When using `DatabaseSessionService` with MySQL backend in google-adk v1.17.0, the application crashes with the following error: app.resources.runner:event_generator:260 - Error in event_generator: (builtins.TypeError) 'tuple' object cannot be interpreted as an integer
<img width="1237" height="129" alt="image" src="https://github.com/user-attachments/assets/0a5fc223-600a-4a92-8443-4d37fb1267f6" />

Root cause: The `DynamicPickleType` class in `database_session_service.py` configures MySQL dialect to use `LONGBLOB` for storing pickled data (line 117-118), but the `process_bind_param` and `process_result_value` methods only handle pickle serialization/deserialization for Spanner dialect, not MySQL. This causes MySQL to attempt storing raw Python objects instead of pickled bytes, leading to serialization errors and potential data corruption.

**Solution:**
Added MySQL to the pickle serialization logic in both `process_bind_param` and `process_result_value` methods, treating it the same way as Spanner dialect. This ensures that:
- Data is properly pickled to bytes before being stored in MySQL's LONGBLOB column
- Data is properly unpickled when retrieved from the database
- No breaking changes to existing functionality for other dialects (SQLite, PostgreSQL)

### Testing Plan

_Please describe the tests that you ran to verify your changes. This is required
for all PRs that are not small documentation or typo fixes._

**Unit Tests:**

- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.

**Summary of `pytest` results:**
<img width="929" height="306" alt="image" src="https://github.com/user-attachments/assets/3d548b96-ac49-4101-8405-a289a722293c" />

**Manual End-to-End (E2E) Tests:**

_Please provide instructions on how to manually test your changes, including any
necessary setup or configuration. Please provide logs or screenshots to help
reviewers better understand the fix._

### Checklist

- [x]  I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x]  I have performed a self-review of my own code.
- [x]  I have commented my code, particularly in hard-to-understand areas.
- [x]  I have added tests that prove my fix is effective or that my feature works.
- [x]  New and existing unit tests pass locally with my changes.
- [x]  I have manually tested my changes end-to-end.
- [x]  Any dependent changes have been merged and published in downstream modules.

### Additional context

_Add any other context or screenshots about the feature request here._

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/3282 from hung12ct:fix/mysql-pickle-serialization d9df37adb7dfbbfd8502a0fe65c4f8bca3d0d978
PiperOrigin-RevId: 825834360
This commit is contained in:
hung12ct
2025-10-29 20:39:58 -07:00
committed by Copybara-Service
parent f9569bbb1a
commit fc15c9a0c3
2 changed files with 183 additions and 2 deletions
@@ -125,14 +125,14 @@ class DynamicPickleType(TypeDecorator):
def process_bind_param(self, value, dialect): def process_bind_param(self, value, dialect):
"""Ensures the pickled value is a bytes object before passing it to the database dialect.""" """Ensures the pickled value is a bytes object before passing it to the database dialect."""
if value is not None: if value is not None:
if dialect.name == "spanner+spanner": if dialect.name in ("spanner+spanner", "mysql"):
return pickle.dumps(value) return pickle.dumps(value)
return value return value
def process_result_value(self, value, dialect): def process_result_value(self, value, dialect):
"""Ensures the raw bytes from the database are unpickled back into a Python object.""" """Ensures the raw bytes from the database are unpickled back into a Python object."""
if value is not None: if value is not None:
if dialect.name == "spanner+spanner": if dialect.name in ("spanner+spanner", "mysql"):
return pickle.loads(value) return pickle.loads(value)
return value return value
@@ -0,0 +1,181 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import pickle
from unittest import mock
from google.adk.sessions.database_session_service import DynamicPickleType
import pytest
from sqlalchemy import create_engine
from sqlalchemy.dialects import mysql
@pytest.fixture
def pickle_type():
"""Fixture for DynamicPickleType instance."""
return DynamicPickleType()
def test_load_dialect_impl_mysql(pickle_type):
"""Test that MySQL dialect uses LONGBLOB."""
# Mock the MySQL dialect
mock_dialect = mock.Mock()
mock_dialect.name = "mysql"
# Mock the return value of type_descriptor
mock_longblob_type = mock.Mock()
mock_dialect.type_descriptor.return_value = mock_longblob_type
impl = pickle_type.load_dialect_impl(mock_dialect)
# Verify type_descriptor was called once with mysql.LONGBLOB
mock_dialect.type_descriptor.assert_called_once_with(mysql.LONGBLOB)
# Verify the return value is what we expect
assert impl == mock_longblob_type
def test_load_dialect_impl_spanner(pickle_type):
"""Test that Spanner dialect uses SpannerPickleType."""
# Mock the spanner dialect
mock_dialect = mock.Mock()
mock_dialect.name = "spanner+spanner"
with mock.patch(
"google.cloud.sqlalchemy_spanner.sqlalchemy_spanner.SpannerPickleType"
) as mock_spanner_type:
pickle_type.load_dialect_impl(mock_dialect)
mock_dialect.type_descriptor.assert_called_once_with(mock_spanner_type)
def test_load_dialect_impl_default(pickle_type):
"""Test that other dialects use default PickleType."""
engine = create_engine("sqlite:///:memory:")
dialect = engine.dialect
impl = pickle_type.load_dialect_impl(dialect)
# Should return the default impl (PickleType)
assert impl == pickle_type.impl
@pytest.mark.parametrize(
"dialect_name",
[
pytest.param("mysql", id="mysql"),
pytest.param("spanner+spanner", id="spanner"),
],
)
def test_process_bind_param_pickle_dialects(pickle_type, dialect_name):
"""Test that MySQL and Spanner dialects pickle the value."""
mock_dialect = mock.Mock()
mock_dialect.name = dialect_name
test_data = {"key": "value", "nested": [1, 2, 3]}
result = pickle_type.process_bind_param(test_data, mock_dialect)
# Should be pickled bytes
assert isinstance(result, bytes)
# Should be able to unpickle back to original
assert pickle.loads(result) == test_data
def test_process_bind_param_default(pickle_type):
"""Test that other dialects return value as-is."""
mock_dialect = mock.Mock()
mock_dialect.name = "sqlite"
test_data = {"key": "value"}
result = pickle_type.process_bind_param(test_data, mock_dialect)
# Should return value unchanged (SQLAlchemy's PickleType handles it)
assert result == test_data
def test_process_bind_param_none(pickle_type):
"""Test that None values are handled correctly."""
mock_dialect = mock.Mock()
mock_dialect.name = "mysql"
result = pickle_type.process_bind_param(None, mock_dialect)
assert result is None
@pytest.mark.parametrize(
"dialect_name",
[
pytest.param("mysql", id="mysql"),
pytest.param("spanner+spanner", id="spanner"),
],
)
def test_process_result_value_pickle_dialects(pickle_type, dialect_name):
"""Test that MySQL and Spanner dialects unpickle the value."""
mock_dialect = mock.Mock()
mock_dialect.name = dialect_name
test_data = {"key": "value", "nested": [1, 2, 3]}
pickled_data = pickle.dumps(test_data)
result = pickle_type.process_result_value(pickled_data, mock_dialect)
# Should be unpickled back to original
assert result == test_data
def test_process_result_value_default(pickle_type):
"""Test that other dialects return value as-is."""
mock_dialect = mock.Mock()
mock_dialect.name = "sqlite"
test_data = {"key": "value"}
result = pickle_type.process_result_value(test_data, mock_dialect)
# Should return value unchanged (SQLAlchemy's PickleType handles it)
assert result == test_data
def test_process_result_value_none(pickle_type):
"""Test that None values are handled correctly."""
mock_dialect = mock.Mock()
mock_dialect.name = "mysql"
result = pickle_type.process_result_value(None, mock_dialect)
assert result is None
@pytest.mark.parametrize(
"dialect_name",
[
pytest.param("mysql", id="mysql"),
pytest.param("spanner+spanner", id="spanner"),
],
)
def test_roundtrip_pickle_dialects(pickle_type, dialect_name):
"""Test full roundtrip for MySQL and Spanner: bind -> result."""
mock_dialect = mock.Mock()
mock_dialect.name = dialect_name
original_data = {
"string": "test",
"number": 42,
"list": [1, 2, 3],
"nested": {"a": 1, "b": 2},
}
# Simulate bind (Python -> DB)
bound_value = pickle_type.process_bind_param(original_data, mock_dialect)
assert isinstance(bound_value, bytes)
# Simulate result (DB -> Python)
result_value = pickle_type.process_result_value(bound_value, mock_dialect)
assert result_value == original_data