diff --git a/src/google/adk/sessions/schemas/v1.py b/src/google/adk/sessions/schemas/v1.py index 7ee8c811fd..d093c5b5fd 100644 --- a/src/google/adk/sessions/schemas/v1.py +++ b/src/google/adk/sessions/schemas/v1.py @@ -28,8 +28,10 @@ from typing import Any from google.adk.platform import uuid as platform_uuid +from sqlalchemy import desc from sqlalchemy import ForeignKeyConstraint from sqlalchemy import func +from sqlalchemy import Index from sqlalchemy import inspect from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.orm import DeclarativeBase @@ -187,6 +189,13 @@ class StorageEvent(Base): ["sessions.app_name", "sessions.user_id", "sessions.id"], ondelete="CASCADE", ), + Index( + "idx_events_app_user_session_ts", + "app_name", + "user_id", + "session_id", + desc("timestamp"), + ), ) @classmethod diff --git a/tests/unittests/sessions/migration/test_database_schema.py b/tests/unittests/sessions/migration/test_database_schema.py index ceb5420b14..0781b37861 100644 --- a/tests/unittests/sessions/migration/test_database_schema.py +++ b/tests/unittests/sessions/migration/test_database_schema.py @@ -72,6 +72,13 @@ def get_schema_version(sync_conn): event_col_names = {c['name'] for c in event_cols} assert 'event_data' in event_col_names assert 'actions' not in event_col_names + + # Verify composite index on events table + event_indexes = await conn.run_sync( + lambda sync_conn: inspect(sync_conn).get_indexes('events') + ) + idx_names = {idx['name'] for idx in event_indexes} + assert 'idx_events_app_user_session_ts' in idx_names await engine.dispose()