mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
f9569bbb1a
commit
fc15c9a0c3
@@ -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
|
||||||
Reference in New Issue
Block a user