google--adk-python/tests/unittests/sessions/migration/test_migration.py

# Copyright 2026 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
import os
import pickle

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.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


class TestToSyncUrl:
  """Tests for the to_sync_url function."""

  @pytest.mark.parametrize(
      "input_url,expected_url",
      [
          # PostgreSQL async drivers
          (
              "postgresql+asyncpg://localhost/mydb",
              "postgresql://localhost/mydb",
          ),
          (
              "postgresql+asyncpg://user:pass@localhost:5432/mydb",
              "postgresql://user:pass@localhost:5432/mydb",
          ),
          # PostgreSQL sync drivers (should still strip)
          (
              "postgresql+psycopg2://localhost/mydb",
              "postgresql://localhost/mydb",
          ),
          # MySQL async drivers
          (
              "mysql+aiomysql://localhost/mydb",
              "mysql://localhost/mydb",
          ),
          (
              "mysql+asyncmy://user:pass@localhost:3306/mydb",
              "mysql://user:pass@localhost:3306/mydb",
          ),
          # SQLite async driver
          (
              "sqlite+aiosqlite:///path/to/db.sqlite",
              "sqlite:///path/to/db.sqlite",
          ),
          (
              "sqlite+aiosqlite:///:memory:",
              "sqlite:///:memory:",
          ),
          # URLs without driver specification (unchanged)
          (
              "postgresql://localhost/mydb",
              "postgresql://localhost/mydb",
          ),
          (
              "mysql://localhost/mydb",
              "mysql://localhost/mydb",
          ),
          (
              "sqlite:///path/to/db.sqlite",
              "sqlite:///path/to/db.sqlite",
          ),
          # Edge cases
          (
              "sqlite:///:memory:",
              "sqlite:///:memory:",
          ),
          # Complex URL with query parameters
          (
              "postgresql+asyncpg://user:pass@host/db?ssl=require",
              "postgresql://user:pass@host/db?ssl=require",
          ),
      ],
  )
  def test_to_sync_url(self, input_url, expected_url):
    """Test that async driver specifications are correctly removed."""
    assert _schema_check_utils.to_sync_url(input_url) == expected_url

  def test_to_sync_url_no_scheme_separator(self):
    """Test that URLs without :// are returned unchanged."""
    # This is an invalid URL but the function should handle it gracefully
    assert _schema_check_utils.to_sync_url("not-a-url") == "not-a-url"

  def test_to_sync_url_empty_string(self):
    """Test that empty string is returned unchanged."""
    assert _schema_check_utils.to_sync_url("") == ""


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}"

  # Set up source DB with old pickle schema
  source_engine = create_engine(source_db_url)
  v0.Base.metadata.create_all(source_engine)
  SourceSession = sessionmaker(bind=source_engine)
  source_session = SourceSession()

  # Populate source data
  now = datetime.now(timezone.utc)
  app_state = v0.StorageAppState(
      app_name="app1", state={"akey": 1}, update_time=now
  )
  user_state = v0.StorageUserState(
      app_name="app1", user_id="user1", state={"ukey": 2}, update_time=now
  )
  session = v0.StorageSession(
      app_name="app1",
      user_id="user1",
      id="session1",
      state={"skey": 3},
      create_time=now,
      update_time=now,
  )
  event = v0.StorageEvent(
      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(v1.StorageMetadata).first()
  assert metadata is not None
  assert metadata.key == _schema_check_utils.SCHEMA_VERSION_KEY
  assert metadata.value == _schema_check_utils.SCHEMA_VERSION_1_JSON

  app_state_res = dest_session.query(v1.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(v1.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(v1.StorageSession).first()
  assert session_res is not None
  assert session_res.id == "session1"
  assert session_res.state == {"skey": 3}

  event_res = dest_session.query(v1.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()


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).

  Users often provide async driver URLs (e.g., postgresql+asyncpg://) since
  that's what ADK requires at runtime. The migration tool should handle these
  by automatically converting them to sync URLs.
  """
  source_db_path = tmp_path / "source_pickle_async.db"
  dest_db_path = tmp_path / "dest_json_async.db"
  # Use async driver URLs like users would typically provide
  source_db_url = f"sqlite+aiosqlite:///{source_db_path}"
  dest_db_url = f"sqlite+aiosqlite:///{dest_db_path}"

  # Set up source DB with old pickle schema using sync URL
  sync_source_url = f"sqlite:///{source_db_path}"
  source_engine = create_engine(sync_source_url)
  v0.Base.metadata.create_all(source_engine)
  SourceSession = sessionmaker(bind=source_engine)
  source_session = SourceSession()

  # Populate source data
  now = datetime.now(timezone.utc)
  app_state = v0.StorageAppState(
      app_name="async_app", state={"key": "value"}, update_time=now
  )
  session = v0.StorageSession(
      app_name="async_app",
      user_id="async_user",
      id="async_session",
      state={},
      create_time=now,
      update_time=now,
  )
  source_session.add_all([app_state, session])
  source_session.commit()
  source_session.close()

  # This should NOT raise an error about async drivers (the fix for #4176)
  mfsp.migrate(source_db_url, dest_db_url)

  # Verify destination DB
  sync_dest_url = f"sqlite:///{dest_db_path}"
  dest_engine = create_engine(sync_dest_url)
  DestSession = sessionmaker(bind=dest_engine)
  dest_session = DestSession()

  metadata = dest_session.query(v1.StorageMetadata).first()
  assert metadata is not None
  assert metadata.key == _schema_check_utils.SCHEMA_VERSION_KEY
  assert metadata.value == _schema_check_utils.SCHEMA_VERSION_1_JSON

  app_state_res = dest_session.query(v1.StorageAppState).first()
  assert app_state_res is not None
  assert app_state_res.app_name == "async_app"
  assert app_state_res.state == {"key": "value"}

  session_res = dest_session.query(v1.StorageSession).first()
  assert session_res is not None
  assert session_res.id == "async_session"

  dest_session.close()


def _assert_update_timestamp_tz_is_utc_timestamp(schema_module) -> None:
  engine = create_engine("sqlite:///:memory:")
  schema_module.Base.metadata.create_all(engine)
  SessionLocal = sessionmaker(bind=engine)

  update_time = datetime(2026, 1, 1, 0, 0, 0)
  storage_session = schema_module.StorageSession(
      app_name="app",
      user_id="user",
      id="sid",
      state={},
      create_time=update_time,
      update_time=update_time,
  )

  with SessionLocal() as db:
    db.add(storage_session)
    db.commit()

    fetched = db.get(schema_module.StorageSession, ("app", "user", "sid"))
    assert fetched is not None
    assert isinstance(fetched.update_timestamp_tz, float)
    assert (
        fetched.update_timestamp_tz
        == update_time.replace(tzinfo=timezone.utc).timestamp()
    )


def test_v1_storage_session_update_timestamp_tz() -> None:
  _assert_update_timestamp_tz_is_utc_timestamp(v1)


def test_v0_storage_session_update_timestamp_tz() -> None:
  _assert_update_timestamp_tz_is_utc_timestamp(v0)