feat: support asyncio in session database service

pyproject.toml

@@ -106,6 +106,7 @@ eval = [
 test = [
   # go/keep-sorted start
   "a2a-sdk>=0.3.0,<0.4.0;python_version>='3.10'",
+  "aiosqlite>=0.21.0",
   "anthropic>=0.43.0",                            # For anthropic model tests
   "langchain-community>=0.3.17",
   "langgraph>=0.2.60, <= 0.4.10",                 # For LangGraphAgent

src/google/adk/sessions/database_session_service.py

@@ -29,20 +29,21 @@
 from sqlalchemy import event
 from sqlalchemy import ForeignKeyConstraint
 from sqlalchemy import func
+from sqlalchemy import select
 from sqlalchemy import Text
 from sqlalchemy.dialects import mysql
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.engine import create_engine
-from sqlalchemy.engine import Engine
 from sqlalchemy.exc import ArgumentError
+from sqlalchemy.ext.asyncio import async_sessionmaker
+from sqlalchemy.ext.asyncio import AsyncEngine
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.ext.mutable import MutableDict
 from sqlalchemy.inspection import inspect
 from sqlalchemy.orm import DeclarativeBase
 from sqlalchemy.orm import Mapped
 from sqlalchemy.orm import mapped_column
 from sqlalchemy.orm import relationship
-from sqlalchemy.orm import Session as DatabaseSessionFactory
-from sqlalchemy.orm import sessionmaker
 from sqlalchemy.schema import MetaData
 from sqlalchemy.types import DateTime
 from sqlalchemy.types import PickleType
@@ -385,17 +386,19 @@ class DatabaseSessionService(BaseSessionService):
   """A session service that uses a database for storage."""
 
   def __init__(self, db_url: str, **kwargs: Any):
-    """Initializes the database session service with a database URL."""
+    """
+    Initializes the database session service with a database URL.
+
+    To initialize tables use `initialize_tables` method.
+    """
     # 1. Create DB engine for db connection
-    # 2. Create all tables based on schema
-    # 3. Initialize all properties
+    # 2. Initialize all properties
     try:
-      db_engine = create_engine(db_url, **kwargs)
+      db_engine = create_async_engine(db_url, **kwargs)
 
       if db_engine.dialect.name == "sqlite":
         # Set sqlite pragma to enable foreign keys constraints
-        event.listen(db_engine, "connect", set_sqlite_pragma)
-
+        event.listen(db_engine.sync_engine, "connect", set_sqlite_pragma)
     except Exception as e:
       if isinstance(e, ArgumentError):
         raise ValueError(
@@ -413,18 +416,22 @@ def __init__(self, db_url: str, **kwargs: Any):
     local_timezone = get_localzone()
     logger.info("Local timezone: %s", local_timezone)
 
-    self.db_engine: Engine = db_engine
+    self.db_engine: AsyncEngine = db_engine
     self.metadata: MetaData = MetaData()
-    self.inspector = inspect(self.db_engine)
 
     # DB session factory method
-    self.database_session_factory: sessionmaker[DatabaseSessionFactory] = (
-        sessionmaker(bind=self.db_engine)
+    self.database_session_factory: async_sessionmaker[AsyncSession] = (
+        async_sessionmaker(bind=self.db_engine)
     )
 
-    # Uncomment to recreate DB every time
-    # Base.metadata.drop_all(self.db_engine)
-    Base.metadata.create_all(self.db_engine)
+  async def initialize_tables(self):
+    """Initialize database tables. Call this after creating the service."""
+    async with self.db_engine.begin() as conn:
+      await conn.run_sync(Base.metadata.create_all)
+
+  async def close(self):
+    """Close the database engine and clean up connections."""
+    await self.db_engine.dispose()
 
   @override
   async def create_session(
@@ -441,11 +448,11 @@ async def create_session(
     # 4. Build the session object with generated id
     # 5. Return the session
 
-    with self.database_session_factory() as sql_session:
+    async with self.database_session_factory() as sql_session:
 
       # Fetch app and user states from storage
-      storage_app_state = sql_session.get(StorageAppState, (app_name))
-      storage_user_state = sql_session.get(
+      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_user_state = await sql_session.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -485,9 +492,9 @@ async def create_session(
           state=session_state,
       )
       sql_session.add(storage_session)
-      sql_session.commit()
+      await sql_session.commit()
 
-      sql_session.refresh(storage_session)
+      await sql_session.refresh(storage_session)
 
       # Merge states for response
       merged_state = _merge_state(app_state, user_state, session_state)
@@ -506,8 +513,8 @@ async def get_session(
     # 1. Get the storage session entry from session table
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
-    with self.database_session_factory() as sql_session:
-      storage_session = sql_session.get(
+    async with self.database_session_factory() as sql_session:
+      storage_session = await sql_session.get(
           StorageSession, (app_name, user_id, session_id)
       )
       if storage_session is None:
@@ -519,8 +526,8 @@ async def get_session(
       else:
         timestamp_filter = True
 
-      storage_events = (
-          sql_session.query(StorageEvent)
+      stmt = (
+          select(StorageEvent)
           .filter(StorageEvent.app_name == app_name)
           .filter(StorageEvent.session_id == storage_session.id)
           .filter(StorageEvent.user_id == user_id)
@@ -531,12 +538,13 @@ async def get_session(
               if config and config.num_recent_events
               else None
           )
-          .all()
       )
+      result = await sql_session.execute(stmt)
+      storage_events = result.scalars().all()
 
       # Fetch states from storage
-      storage_app_state = sql_session.get(StorageAppState, (app_name))
-      storage_user_state = sql_session.get(
+      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_user_state = await sql_session.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -557,16 +565,17 @@ async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
     with self.database_session_factory() as sql_session:
-      results = (
-          sql_session.query(StorageSession)
+      stmt = (
+          select(StorageSession)
           .filter(StorageSession.app_name == app_name)
           .filter(StorageSession.user_id == user_id)
-          .all()
       )
+      result = await sql_session.execute(stmt)
+      results = result.scalars().all()
 
       # Fetch states from storage
-      storage_app_state = sql_session.get(StorageAppState, (app_name))
-      storage_user_state = sql_session.get(
+      storage_app_state = await sql_session.get(StorageAppState, (app_name))
+      storage_user_state = await sql_session.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -585,14 +594,14 @@ async def list_sessions(
   async def delete_session(
       self, app_name: str, user_id: str, session_id: str
   ) -> None:
-    with self.database_session_factory() as sql_session:
+    async with self.database_session_factory() as sql_session:
       stmt = delete(StorageSession).where(
           StorageSession.app_name == app_name,
           StorageSession.user_id == user_id,
           StorageSession.id == session_id,
       )
-      sql_session.execute(stmt)
-      sql_session.commit()
+      await sql_session.execute(stmt)
+      await sql_session.commit()
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
@@ -603,7 +612,7 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # 2. Update session attributes based on event config
     # 3. Store event to table
     with self.database_session_factory() as sql_session:
-      storage_session = sql_session.get(
+      storage_session = await sql_session.get(
           StorageSession, (session.app_name, session.user_id, session.id)
       )
 
@@ -617,8 +626,10 @@ async def append_event(self, session: Session, event: Event) -> Event:
         )
 
       # Fetch states from storage
-      storage_app_state = sql_session.get(StorageAppState, (session.app_name))
-      storage_user_state = sql_session.get(
+      storage_app_state = await sql_session.get(
+          StorageAppState, (session.app_name)
+      )
+      storage_user_state = await sql_session.get(
           StorageUserState, (session.app_name, session.user_id)
       )
 
@@ -649,8 +660,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
       sql_session.add(StorageEvent.from_event(session, event))
 
-      sql_session.commit()
-      sql_session.refresh(storage_session)
+      await sql_session.commit()
+      await sql_session.refresh(storage_session)
 
       # Update timestamp with commit time
       session.last_update_time = storage_session.update_timestamp_tz

tests/unittests/sessions/test_session_service.py

@@ -30,12 +30,13 @@ class SessionServiceType(enum.Enum):
   DATABASE = 'DATABASE'
 
 
-def get_session_service(
+async def get_session_service(
     service_type: SessionServiceType = SessionServiceType.IN_MEMORY,
 ):
   """Creates a session service for testing."""
   if service_type == SessionServiceType.DATABASE:
-    return DatabaseSessionService('sqlite:///:memory:')
+    service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+    await service.initialize_tables()
   return InMemorySessionService()
 
 
@@ -44,7 +45,7 @@ def get_session_service(
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_get_empty_session(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   assert not await session_service.get_session(
       app_name='my_app', user_id='test_user', session_id='123'
   )
@@ -55,7 +56,7 @@ async def test_get_empty_session(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_create_get_session(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'test_user'
   state = {'key': 'value'}
@@ -99,7 +100,7 @@ async def test_create_get_session(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_create_and_list_sessions(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'test_user'
 
@@ -126,7 +127,7 @@ async def test_create_and_list_sessions(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_session_state(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id_1 = 'user1'
   user_id_2 = 'user2'
@@ -218,7 +219,7 @@ async def test_session_state(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_create_new_session_will_merge_states(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'user'
   session_id_1 = 'session1'
@@ -264,7 +265,7 @@ async def test_create_new_session_will_merge_states(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_append_event_bytes(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'user'
 
@@ -305,7 +306,7 @@ async def test_append_event_bytes(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_append_event_complete(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'user'
 
@@ -345,7 +346,7 @@ async def test_append_event_complete(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_get_session_with_config(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'user'
 
@@ -409,7 +410,7 @@ async def test_get_session_with_config(service_type):
     'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
 )
 async def test_append_event_with_fields(service_type):
-  session_service = get_session_service(service_type)
+  session_service = await get_session_service(service_type)
   app_name = 'my_app'
   user_id = 'test_user'
   session = await session_service.create_session(