google--adk-python/tests/unittests/sessions/migration/test_database_schema.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.

from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.migration import _schema_check_utils
from google.adk.sessions.schemas import v0
import pytest
from sqlalchemy import inspect
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine


async def create_v0_db(db_path):
  db_url = f'sqlite+aiosqlite:///{db_path}'
  engine = create_async_engine(db_url)
  async with engine.begin() as conn:
    await conn.run_sync(v0.Base.metadata.create_all)
  await engine.dispose()


# Use async context managers so DatabaseSessionService always closes.


@pytest.mark.asyncio
async def test_new_db_uses_latest_schema(tmp_path):
  db_path = tmp_path / 'new_db.db'
  db_url = f'sqlite+aiosqlite:///{db_path}'
  async with DatabaseSessionService(db_url) as session_service:
    assert session_service._db_schema_version is None
    await session_service.create_session(app_name='my_app', user_id='test_user')
    assert (
        session_service._db_schema_version
        == _schema_check_utils.LATEST_SCHEMA_VERSION
    )

  # Verify metadata table
  engine = create_async_engine(db_url)
  async with engine.connect() as conn:
    has_metadata_table = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
    )
    assert has_metadata_table

    def get_schema_version(sync_conn):
      inspector = inspect(sync_conn)
      key_col = inspector.dialect.identifier_preparer.quote('key')
      return sync_conn.execute(
          text(
              f'SELECT value FROM adk_internal_metadata WHERE {key_col} = :key'
          ),
          {'key': _schema_check_utils.SCHEMA_VERSION_KEY},
      ).scalar_one_or_none()

    schema_version = await conn.run_sync(get_schema_version)
    assert schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION

    # Verify events table columns for v1
    event_cols = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_columns('events')
    )
    event_col_names = {c['name'] for c in event_cols}
    assert 'event_data' in event_col_names
    assert 'actions' not in event_col_names

    event_indexes = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_indexes('events')
    )
    assert any(
        index['name'] == 'idx_events_app_user_session_ts'
        and index['column_names']
        == ['app_name', 'user_id', 'session_id', 'timestamp']
        for index in event_indexes
    )
  await engine.dispose()


@pytest.mark.asyncio
async def test_existing_v0_db_uses_v0_schema(tmp_path):
  db_path = tmp_path / 'v0_db.db'
  await create_v0_db(db_path)
  db_url = f'sqlite+aiosqlite:///{db_path}'
  async with DatabaseSessionService(db_url) as session_service:
    assert session_service._db_schema_version is None
    await session_service.create_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    assert (
        session_service._db_schema_version
        == _schema_check_utils.SCHEMA_VERSION_0_PICKLE
    )

    session = await session_service.get_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    assert session.id == 's1'

  # Verify schema tables
  engine = create_async_engine(db_url)
  async with engine.connect() as conn:
    has_metadata_table = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
    )
    assert not has_metadata_table

    # Verify events table columns for v0
    event_cols = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_columns('events')
    )
    event_col_names = {c['name'] for c in event_cols}
    assert 'event_data' not in event_col_names
    assert 'actions' in event_col_names
  await engine.dispose()


@pytest.mark.asyncio
async def test_existing_latest_db_uses_latest_schema(tmp_path):
  db_path = tmp_path / 'new_db.db'
  db_url = f'sqlite+aiosqlite:///{db_path}'

  # Create session service which creates db with latest schema
  async with DatabaseSessionService(db_url) as session_service1:
    await session_service1.create_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    assert (
        session_service1._db_schema_version
        == _schema_check_utils.LATEST_SCHEMA_VERSION
    )

    # Create another session service on same db and check it detects latest schema
    async with DatabaseSessionService(db_url) as session_service2:
      await session_service2.create_session(
          app_name='my_app', user_id='test_user2', session_id='s2'
      )
      assert (
          session_service2._db_schema_version
          == _schema_check_utils.LATEST_SCHEMA_VERSION
      )
      s2 = await session_service2.get_session(
          app_name='my_app', user_id='test_user2', session_id='s2'
      )
      assert s2.id == 's2'

      s1 = await session_service2.get_session(
          app_name='my_app', user_id='test_user', session_id='s1'
      )
      assert s1.id == 's1'

      list_sessions_response = await session_service2.list_sessions(
          app_name='my_app'
      )
      assert len(list_sessions_response.sessions) == 2

  # Verify schema tables
  engine = create_async_engine(db_url)
  async with engine.connect() as conn:
    has_metadata_table = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).has_table('adk_internal_metadata')
    )
    assert has_metadata_table

    # Verify events table columns for v1
    event_cols = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_columns('events')
    )
    event_col_names = {c['name'] for c in event_cols}
    assert 'event_data' in event_col_names
    assert 'actions' not in event_col_names
  await engine.dispose()


@pytest.mark.asyncio
async def test_prepare_tables_recreates_missing_latest_events_index(tmp_path):
  db_path = tmp_path / 'missing_latest_index.db'
  db_url = f'sqlite+aiosqlite:///{db_path}'

  async with DatabaseSessionService(db_url) as session_service:
    await session_service.create_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )

  engine = create_async_engine(db_url)
  async with engine.begin() as conn:
    await conn.execute(text('DROP INDEX idx_events_app_user_session_ts'))
  await engine.dispose()

  async with DatabaseSessionService(db_url) as session_service:
    session = await session_service.get_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    assert session.id == 's1'

  engine = create_async_engine(db_url)
  async with engine.connect() as conn:
    event_indexes = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_indexes('events')
    )
  await engine.dispose()

  assert any(
      index['name'] == 'idx_events_app_user_session_ts'
      and index['column_names']
      == ['app_name', 'user_id', 'session_id', 'timestamp']
      for index in event_indexes
  )


@pytest.mark.asyncio
async def test_prepare_tables_recreates_missing_v0_events_index(tmp_path):
  db_path = tmp_path / 'missing_v0_index.db'
  await create_v0_db(db_path)
  db_url = f'sqlite+aiosqlite:///{db_path}'

  engine = create_async_engine(db_url)
  async with engine.begin() as conn:
    await conn.execute(text('DROP INDEX idx_events_app_user_session_ts'))
  await engine.dispose()

  async with DatabaseSessionService(db_url) as session_service:
    await session_service.create_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    session = await session_service.get_session(
        app_name='my_app', user_id='test_user', session_id='s1'
    )
    assert session.id == 's1'

  engine = create_async_engine(db_url)
  async with engine.connect() as conn:
    event_indexes = await conn.run_sync(
        lambda sync_conn: inspect(sync_conn).get_indexes('events')
    )
  await engine.dispose()

  assert any(
      index['name'] == 'idx_events_app_user_session_ts'
      and index['column_names']
      == ['app_name', 'user_id', 'session_id', 'timestamp']
      for index in event_indexes
  )