chore: Close database engines to avoid aiosqlite pytest hangs

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 852428755

Author: GWeale

pyproject.toml

@@ -26,8 +26,7 @@ classifiers = [ # List of https://pypi.org/classifiers/
 dependencies = [
   # go/keep-sorted start
   "PyYAML>=6.0.2, <7.0.0",                                   # For APIHubToolset.
-  # TODO: Update aiosqlite version once https://github.com/omnilib/aiosqlite/issues/369 is fixed.
-  "aiosqlite==0.21.0",                                       # For SQLite database
+  "aiosqlite>=0.21.0",                                       # For SQLite database
   "anyio>=4.9.0, <5.0.0",                                    # For MCP Session Manager
   "authlib>=1.5.1, <2.0.0",                                  # For RestAPI Tool
   "click>=8.1.8, <9.0.0",                                    # For CLI tools
@@ -110,6 +109,7 @@ eval = [
   "google-cloud-aiplatform[evaluation]>=1.100.0",
   "pandas>=2.2.3",
   "rouge-score>=0.1.2",
+  "scipy<1.16; python_version<'3.11'",
   "tabulate>=0.9.0",
   # go/keep-sorted end
 ]

src/google/adk/sessions/database_session_service.py

@@ -25,12 +25,14 @@
 from sqlalchemy import event
 from sqlalchemy import select
 from sqlalchemy import text
+from sqlalchemy.engine import make_url
 from sqlalchemy.exc import ArgumentError
 from sqlalchemy.ext.asyncio import async_sessionmaker
 from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.asyncio import AsyncSession as DatabaseSessionFactory
 from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.inspection import inspect
+from sqlalchemy.pool import StaticPool
 from typing_extensions import override
 from tzlocal import get_localzone
 
@@ -103,7 +105,15 @@ def __init__(self, db_url: str, **kwargs: Any):
     # 2. Create all tables based on schema
     # 3. Initialize all properties
     try:
-      db_engine = create_async_engine(db_url, **kwargs)
+      engine_kwargs = dict(kwargs)
+      url = make_url(db_url)
+      if url.get_backend_name() == "sqlite" and url.database == ":memory:":
+        engine_kwargs.setdefault("poolclass", StaticPool)
+        connect_args = dict(engine_kwargs.get("connect_args", {}))
+        connect_args.setdefault("check_same_thread", False)
+        engine_kwargs["connect_args"] = connect_args
+
+      db_engine = create_async_engine(db_url, **engine_kwargs)
       if db_engine.dialect.name == "sqlite":
         # Set sqlite pragma to enable foreign keys constraints
         event.listen(db_engine.sync_engine, "connect", _set_sqlite_pragma)
@@ -477,3 +487,15 @@ async def append_event(self, session: Session, event: Event) -> Event:
     # Also update the in-memory session
     await super().append_event(session=session, event=event)
     return event
+
+  async def close(self) -> None:
+    """Disposes the SQLAlchemy engine and closes pooled connections."""
+    await self.db_engine.dispose()
+
+  async def __aenter__(self) -> DatabaseSessionService:
+    """Enters the async context manager and returns this service."""
+    return self
+
+  async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+    """Exits the async context manager and closes the service."""
+    await self.close()

tests/unittests/sessions/migration/test_database_schema.py

@@ -29,17 +29,20 @@ async def create_v0_db(db_path):
   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}'
-  session_service = DatabaseSessionService(db_url)
-  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
-  )
+  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)
@@ -71,21 +74,20 @@ 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}'
-  session_service = DatabaseSessionService(db_url)
-
-  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'
+  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)
@@ -111,38 +113,38 @@ async def test_existing_latest_db_uses_latest_schema(tmp_path):
   db_url = f'sqlite+aiosqlite:///{db_path}'
 
   # Create session service which creates db with latest schema
-  session_service1 = DatabaseSessionService(db_url)
-  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
-  session_service2 = DatabaseSessionService(db_url)
-  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
+  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)