mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat!: Rollback the DB migration as it is breaking
Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 839818479
This commit is contained in:
Shangjie Chen
committed by
Copybara-Service
Copybara-Service
parent
8c9105bf14
commit
9d918d45df
@@ -1,106 +0,0 @@
|
||||
# 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.
|
||||
"""Tests for migration scripts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from datetime import timezone
|
||||
|
||||
from google.adk.events.event_actions import EventActions
|
||||
from google.adk.sessions import database_session_service as dss
|
||||
from google.adk.sessions.migration import _schema_check
|
||||
from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle as mfsp
|
||||
import pytest
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
|
||||
def test_migrate_from_sqlalchemy_pickle(tmp_path):
|
||||
"""Tests for migrate_from_sqlalchemy_pickle."""
|
||||
source_db_path = tmp_path / "source_pickle.db"
|
||||
dest_db_path = tmp_path / "dest_json.db"
|
||||
source_db_url = f"sqlite:///{source_db_path}"
|
||||
dest_db_url = f"sqlite:///{dest_db_path}"
|
||||
|
||||
# Setup source DB with old pickle schema
|
||||
source_engine = create_engine(source_db_url)
|
||||
mfsp.OldBase.metadata.create_all(source_engine)
|
||||
SourceSession = sessionmaker(bind=source_engine)
|
||||
source_session = SourceSession()
|
||||
|
||||
# Populate source data
|
||||
now = datetime.now(timezone.utc)
|
||||
app_state = mfsp.OldStorageAppState(
|
||||
app_name="app1", state={"akey": 1}, update_time=now
|
||||
)
|
||||
user_state = mfsp.OldStorageUserState(
|
||||
app_name="app1", user_id="user1", state={"ukey": 2}, update_time=now
|
||||
)
|
||||
session = mfsp.OldStorageSession(
|
||||
app_name="app1",
|
||||
user_id="user1",
|
||||
id="session1",
|
||||
state={"skey": 3},
|
||||
create_time=now,
|
||||
update_time=now,
|
||||
)
|
||||
event = mfsp.OldStorageEvent(
|
||||
id="event1",
|
||||
app_name="app1",
|
||||
user_id="user1",
|
||||
session_id="session1",
|
||||
invocation_id="invoke1",
|
||||
author="user",
|
||||
actions=EventActions(state_delta={"skey": 4}),
|
||||
timestamp=now,
|
||||
)
|
||||
source_session.add_all([app_state, user_state, session, event])
|
||||
source_session.commit()
|
||||
source_session.close()
|
||||
|
||||
mfsp.migrate(source_db_url, dest_db_url)
|
||||
|
||||
# Verify destination DB
|
||||
dest_engine = create_engine(dest_db_url)
|
||||
DestSession = sessionmaker(bind=dest_engine)
|
||||
dest_session = DestSession()
|
||||
|
||||
metadata = dest_session.query(dss.StorageMetadata).first()
|
||||
assert metadata is not None
|
||||
assert metadata.key == _schema_check.SCHEMA_VERSION_KEY
|
||||
assert metadata.value == _schema_check.SCHEMA_VERSION_1_0_JSON
|
||||
|
||||
app_state_res = dest_session.query(dss.StorageAppState).first()
|
||||
assert app_state_res is not None
|
||||
assert app_state_res.app_name == "app1"
|
||||
assert app_state_res.state == {"akey": 1}
|
||||
|
||||
user_state_res = dest_session.query(dss.StorageUserState).first()
|
||||
assert user_state_res is not None
|
||||
assert user_state_res.user_id == "user1"
|
||||
assert user_state_res.state == {"ukey": 2}
|
||||
|
||||
session_res = dest_session.query(dss.StorageSession).first()
|
||||
assert session_res is not None
|
||||
assert session_res.id == "session1"
|
||||
assert session_res.state == {"skey": 3}
|
||||
|
||||
event_res = dest_session.query(dss.StorageEvent).first()
|
||||
assert event_res is not None
|
||||
assert event_res.id == "event1"
|
||||
assert "state_delta" in event_res.event_data["actions"]
|
||||
assert event_res.event_data["actions"]["state_delta"] == {"skey": 4}
|
||||
|
||||
dest_session.close()
|
||||
@@ -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