feat: Add support for reusing an existing SQLAlchemy AsyncEngine in DatabaseSessionService

Merges #3565

Also resolved reviewer comments:
- Track engine ownership to avoid disposing caller-owned engine on close.
- Avoid mutating caller-owned engine with sqlite connection pragma listeners.
- Added tests to verify caller-owned engine remains usable after service close.
- Configured manual engine with StaticPool in tests to avoid flakiness.

Co-authored-by: Shangjie Chen <deanchen@google.com>
COPYBARA_INTEGRATE_REVIEW=#3565 from GitMarco27:feature/database-session-service-accept-engine 2abdee6
PiperOrigin-RevId: 934104748

Author: 2 people

src/google/adk/sessions/database_session_service.py

@@ -22,6 +22,7 @@
 from typing import Any
 from typing import AsyncIterator
 from typing import Optional
+from typing import overload
 from typing import TypeAlias
 from typing import TypeVar
 
@@ -193,51 +194,100 @@ def __init__(self, version: str):
 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."""
-    # 1. Create DB engine for db connection
-    # 2. Create all tables based on schema
-    # 3. Initialize all properties
+  @overload
+  def __init__(
+      self,
+      db_url: str,
+      **kwargs: Any,
+  ) -> None:
+    """Initializes the database session service with a database URL.
+
+    Args:
+      db_url: Database URL string for creating a new engine.
+      **kwargs: Additional keyword arguments passed to create_async_engine.
+    """
+
+  @overload
+  def __init__(
+      self,
+      *,
+      db_engine: AsyncEngine,
+  ) -> None:
+    """Initializes the database session service with an existing SQLAlchemy AsyncEngine.
+
+    Args:
+      db_engine: Existing SQLAlchemy AsyncEngine instance to use.
+    """
+
+  def __init__(
+      self,
+      db_url: Optional[str] = None,
+      db_engine: Optional[AsyncEngine] = None,
+      **kwargs: Any,
+  ) -> None:
+    """Initializes the database session service.
+
+    Args:
+      db_url: Database URL string for creating a new engine. Mutually exclusive
+        with db_engine.
+      db_engine: Existing AsyncEngine instance. Mutually exclusive with db_url.
+      **kwargs: Additional keyword arguments passed to create_async_engine when
+        db_url is provided. Ignored when db_engine is provided.
+
+    Raises:
+      ValueError: If neither or both db_url and db_engine are provided, or if
+        engine creation fails.
+    """
     try:
       import sqlalchemy  # noqa: F401
     except ImportError as e:
       from ..utils._dependency import missing_extra
 
       raise missing_extra("sqlalchemy", "db") from e
 
-    try:
-      engine_kwargs = dict(kwargs)
-      url = make_url(db_url)
-      if (
-          url.get_backend_name() == _SQLITE_DIALECT
-          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
-      elif url.get_backend_name() != _SQLITE_DIALECT:
-        engine_kwargs.setdefault("pool_pre_ping", True)
-
-      db_engine = create_async_engine(db_url, **engine_kwargs)
-      if db_engine.dialect.name == _SQLITE_DIALECT:
-        # Set sqlite pragma to enable foreign keys constraints
-        event.listen(db_engine.sync_engine, "connect", _set_sqlite_pragma)
-
-    except Exception as e:
-      if isinstance(e, ArgumentError):
-        raise ValueError(
-            f"Invalid database URL format or argument '{db_url}'."
-        ) from e
-      if isinstance(e, ImportError):
+    if (db_url is None) == (db_engine is None):
+      raise ValueError(
+          "Exactly one of 'db_url' or 'db_engine' must be provided."
+      )
+
+    if db_engine is None:
+      self._owns_db_engine = True
+      try:
+        engine_kwargs = dict(kwargs)
+        url = make_url(db_url)
+        if (
+            url.get_backend_name() == _SQLITE_DIALECT
+            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
+        elif url.get_backend_name() != _SQLITE_DIALECT:
+          engine_kwargs.setdefault("pool_pre_ping", True)
+
+        db_engine = create_async_engine(db_url, **engine_kwargs)
+        if db_engine.dialect.name == _SQLITE_DIALECT:
+          # Set sqlite pragma to enable foreign keys constraints
+          event.listen(db_engine.sync_engine, "connect", _set_sqlite_pragma)
+
+      except Exception as e:
+        if isinstance(e, ArgumentError):
+          raise ValueError(
+              f"Invalid database URL format or argument '{db_url}'."
+          ) from e
+        if isinstance(e, ImportError):
+          raise ValueError(
+              f"Database related module not found for URL '{db_url}'."
+          ) from e
         raise ValueError(
-            f"Database related module not found for URL '{db_url}'."
+            f"Failed to create database engine for URL '{db_url}'"
         ) from e
-      raise ValueError(
-          f"Failed to create database engine for URL '{db_url}'"
-      ) from e
+    else:
+      self._owns_db_engine = False
 
     self.db_engine: AsyncEngine = db_engine
+
     # DB session factory method
     self.database_session_factory: async_sessionmaker[
         DatabaseSessionFactory
@@ -802,7 +852,8 @@ async def append_event(self, session: Session, event: Event) -> Event:
 
   async def close(self) -> None:
     """Disposes the SQLAlchemy engine and closes pooled connections."""
-    await self.db_engine.dispose()
+    if self._owns_db_engine:
+      await self.db_engine.dispose()
 
   async def __aenter__(self) -> DatabaseSessionService:
     """Enters the async context manager and returns this service."""

tests/unittests/sessions/test_session_service.py

@@ -34,6 +34,9 @@
 from google.genai import types
 import pytest
 from sqlalchemy import delete
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.pool import StaticPool
 
 
 class SessionServiceType(enum.Enum):
@@ -1951,3 +1954,108 @@ def test_database_session_service_visible_in_module_namespace():
 
   assert 'DatabaseSessionService' in dir(sessions_module)
   assert sessions_module.DatabaseSessionService is DatabaseSessionService
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_with_db_url():
+  """Test DatabaseSessionService initialization with db_url."""
+  # Test db_url as positional argument
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  app_name = 'test_app'
+  user_id = 'test_user'
+
+  # Create and retrieve a session
+  session = await service.create_session(
+      app_name=app_name, user_id=user_id, state={'key': 'value'}
+  )
+  assert session.app_name == app_name
+  assert session.user_id == user_id
+  assert session.state == {'key': 'value'}
+
+  # Let's check that we can retrieve it
+  retrieved = await service.get_session(
+      app_name=app_name, user_id=user_id, session_id=session.id
+  )
+  assert retrieved == session
+
+  # test db_url as keyword argument
+  service2 = DatabaseSessionService(db_url='sqlite+aiosqlite:///:memory:')
+  session2 = await service2.create_session(
+      app_name=app_name, user_id=user_id, state={'key': 'value2'}
+  )
+  assert session2.state == {'key': 'value2'}
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_with_db_engine():
+  """Test DatabaseSessionService initialization with db_engine."""
+  # Create an engine manually with StaticPool to avoid flakes
+  engine = create_async_engine(
+      'sqlite+aiosqlite:///:memory:',
+      poolclass=StaticPool,
+      connect_args={'check_same_thread': False},
+  )
+
+  # Create service with db_engine
+  service = DatabaseSessionService(db_engine=engine)
+  app_name = 'test_app'
+  user_id = 'test_user'
+
+  # Create and retrieve a session
+  session = await service.create_session(
+      app_name=app_name, user_id=user_id, state={'key': 'value'}
+  )
+  assert session.app_name == app_name
+  assert session.user_id == user_id
+  assert session.state == {'key': 'value'}
+
+  # Let's check that we can retrieve it
+  retrieved = await service.get_session(
+      app_name=app_name, user_id=user_id, session_id=session.id
+  )
+  assert retrieved == session
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_caller_owned_engine_not_disposed_on_close():
+  """Verifies that a caller-owned engine is not disposed when the service is closed."""
+  engine = create_async_engine(
+      'sqlite+aiosqlite:///:memory:',
+      poolclass=StaticPool,
+      connect_args={'check_same_thread': False},
+  )
+
+  service = DatabaseSessionService(db_engine=engine)
+
+  # Use the service
+  session = await service.create_session(app_name='app', user_id='user')
+  assert session is not None
+
+  # Close the service
+  await service.close()
+
+  # Verify engine is still usable by running a query
+  async with engine.connect() as conn:
+    result = await conn.execute(text('SELECT 1;'))
+    assert result.scalar() == 1
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_requires_one_argument():
+  """Test that DatabaseSessionService requires exactly one of db_url or db_engine."""
+  # Neither argument provided
+  with pytest.raises(
+      ValueError,
+      match="Exactly one of 'db_url' or 'db_engine' must be provided",
+  ):
+    DatabaseSessionService()
+
+  # Both arguments provided
+  engine = create_async_engine('sqlite+aiosqlite:///:memory:')
+  with pytest.raises(
+      ValueError,
+      match="Exactly one of 'db_url' or 'db_engine' must be provided",
+  ):
+    DatabaseSessionService(
+        db_url='sqlite+aiosqlite:///:memory:', db_engine=engine
+    )