feat: Add index to events table and update dependencies

GWeale · copybara-github · commit 3153e6d74f40 · 2026-03-24T15:10:58.000-07:00
Introduces a new index `idx_events_app_user_session_ts` on the `events` table in both V0 and V1 database schemas. This index covers `app_name`, `user_id`, `session_id`, and `timestamp` (descending) to improve query performance. Modifies `DatabaseSessionService` to automatically create any missing indexes defined in the SQLAlchemy metadata when preparing the database tables. Close #4827 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 888877174
diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py
@@ -21,13 +21,16 @@
 import logging
 from typing import Any
 from typing import AsyncIterator
+from typing import Optional
 from typing import TypeAlias
 from typing import TypeVar
 
 from google.adk.platform import time as platform_time
 from sqlalchemy import delete
 from sqlalchemy import event
+from sqlalchemy import MetaData
 from sqlalchemy import select
+from sqlalchemy.engine import Connection
 from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
 from sqlalchemy.ext.asyncio import async_sessionmaker
@@ -106,6 +109,22 @@ def _set_sqlite_pragma(dbapi_connection, connection_record):
   cursor.close()
 
 
+def _ensure_schema_indexes_exist(
+    connection: Connection, metadata: MetaData
+) -> None:
+  """Ensures indexes declared in metadata exist for existing tables."""
+  logger.debug("Ensuring schema indexes exist for metadata tables.")
+  for table in metadata.sorted_tables:
+    for index in sorted(table.indexes, key=lambda item: item.name or ""):
+      index.create(bind=connection, checkfirst=True)
+
+
+def _setup_database_schema(connection: Connection, metadata: MetaData) -> None:
+  """Ensures tables and indexes declared in metadata exist."""
+  metadata.create_all(bind=connection)
+  _ensure_schema_indexes_exist(connection, metadata)
+
+
 def _merge_state(
     app_state: dict[str, Any],
     user_state: dict[str, Any],
@@ -301,11 +320,11 @@ async def _prepare_tables(self):
           # Uncomment to recreate DB every time
           # await conn.run_sync(BaseV1.metadata.drop_all)
           logger.debug("Using V1 schema tables...")
-          await conn.run_sync(BaseV1.metadata.create_all)
+          await conn.run_sync(_setup_database_schema, BaseV1.metadata)
         else:
           # await conn.run_sync(BaseV0.metadata.drop_all)
           logger.debug("Using V0 schema tables...")
-          await conn.run_sync(BaseV0.metadata.create_all)
+          await conn.run_sync(_setup_database_schema, BaseV0.metadata)
 
       if self._db_schema_version == _schema_check_utils.LATEST_SCHEMA_VERSION:
         async with self._rollback_on_exception_session() as sql_session:
diff --git a/src/google/adk/sessions/schemas/v0.py b/src/google/adk/sessions/schemas/v0.py
@@ -36,8 +36,10 @@
 from google.adk.platform import uuid as platform_uuid
 from google.genai import types
 from sqlalchemy import Boolean
+from sqlalchemy import desc
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import func
+from sqlalchemy import Index
 from sqlalchemy import inspect
 from sqlalchemy import Text
 from sqlalchemy.dialects import mysql
@@ -256,6 +258,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"),
+      ),
   )
 
   @property
diff --git 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
@@ -196,6 +198,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
@@ -72,6 +72,16 @@ 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
+
+    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()
 
 
@@ -168,3 +178,74 @@ async def test_existing_latest_db_uses_latest_schema(tmp_path):
     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
+  )
