fix: Add read-only session support in DatabaseSessionService

This change introduces a separate async session factory for Spanner connections configured with `read_only=True`

Close #4771

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

Author: GWeale

src/google/adk/sessions/database_session_service.py

@@ -21,7 +21,6 @@
 import logging
 from typing import Any
 from typing import AsyncIterator
-from typing import Optional
 from typing import TypeAlias
 from typing import TypeVar
 
@@ -178,11 +177,14 @@ def __init__(self, db_url: str, **kwargs: Any):
       ) from e
 
     self.db_engine: AsyncEngine = db_engine
-
     # DB session factory method
     self.database_session_factory: async_sessionmaker[
         DatabaseSessionFactory
     ] = async_sessionmaker(bind=self.db_engine, expire_on_commit=False)
+    read_only_engine = self.db_engine.execution_options(read_only=True)
+    self._read_only_database_session_factory: async_sessionmaker[
+        DatabaseSessionFactory
+    ] = async_sessionmaker(bind=read_only_engine, expire_on_commit=False)
 
     # Flag to indicate if tables are created
     self._tables_created = False
@@ -201,17 +203,27 @@ def __init__(self, db_url: str, **kwargs: Any):
   def _get_schema_classes(self) -> _SchemaClasses:
     return _SchemaClasses(self._db_schema_version)
 
+  def _get_database_session_factory(
+      self, *, read_only: bool = False
+  ) -> async_sessionmaker[DatabaseSessionFactory]:
+    if read_only:
+      return self._read_only_database_session_factory
+    return self.database_session_factory
+
   @asynccontextmanager
   async def _rollback_on_exception_session(
       self,
+      *,
+      read_only: bool = False,
   ) -> AsyncIterator[DatabaseSessionFactory]:
     """Yields a database session with guaranteed rollback on errors.
 
     On normal exit the caller is responsible for committing; on any exception
     the transaction is explicitly rolled back before the error propagates,
     preventing connection-pool exhaustion from lingering invalid transactions.
     """
-    async with self.database_session_factory() as sql_session:
+    session_factory = self._get_database_session_factory(read_only=read_only)
+    async with session_factory() as sql_session:
       try:
         yield sql_session
       except BaseException:
@@ -441,7 +453,9 @@ async def get_session(
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
     schema = self._get_schema_classes()
-    async with self._rollback_on_exception_session() as sql_session:
+    async with self._rollback_on_exception_session(
+        read_only=True
+    ) as sql_session:
       storage_session = await sql_session.get(
           schema.StorageSession, (app_name, user_id, session_id)
       )
@@ -496,7 +510,9 @@ async def list_sessions(
   ) -> ListSessionsResponse:
     await self._prepare_tables()
     schema = self._get_schema_classes()
-    async with self._rollback_on_exception_session() as sql_session:
+    async with self._rollback_on_exception_session(
+        read_only=True
+    ) as sql_session:
       stmt = select(schema.StorageSession).filter(
           schema.StorageSession.app_name == app_name
       )

tests/unittests/sessions/test_session_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+from contextlib import asynccontextmanager
 from datetime import datetime
 from datetime import timezone
 import enum
@@ -151,6 +152,74 @@ def fake_create_async_engine(_db_url: str, **kwargs):
   assert captured_kwargs.get('pool_pre_ping') is False
 
 
+def test_database_session_service_creates_read_only_engine_for_spanner():
+  captured_binds = []
+  fake_engine = mock.Mock()
+  fake_engine.dialect.name = 'spanner+spanner'
+  fake_engine.sync_engine = mock.Mock()
+  read_only_engine = mock.Mock()
+  fake_engine.execution_options.return_value = read_only_engine
+
+  def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
+    del expire_on_commit
+    del kwargs
+    captured_binds.append(bind)
+    return mock.Mock()
+
+  with (
+      mock.patch.object(
+          database_session_service,
+          'create_async_engine',
+          return_value=fake_engine,
+      ),
+      mock.patch.object(
+          database_session_service,
+          'async_sessionmaker',
+          side_effect=fake_async_sessionmaker,
+      ),
+  ):
+    database_session_service.DatabaseSessionService(
+        'spanner+spanner:///projects/test/instances/test/databases/test'
+    )
+
+  assert captured_binds == [fake_engine, read_only_engine]
+  fake_engine.execution_options.assert_called_once_with(read_only=True)
+
+
+def test_database_session_service_creates_read_only_engine_for_other_dialects():
+  captured_binds = []
+  fake_engine = mock.Mock()
+  fake_engine.dialect.name = 'postgresql'
+  fake_engine.sync_engine = mock.Mock()
+  read_only_engine = mock.Mock()
+  fake_engine.execution_options.return_value = read_only_engine
+
+  def fake_async_sessionmaker(*, bind, expire_on_commit, **kwargs):
+    del expire_on_commit
+    del kwargs
+    captured_binds.append(bind)
+    return mock.Mock()
+
+  with (
+      mock.patch.object(
+          database_session_service,
+          'create_async_engine',
+          return_value=fake_engine,
+      ),
+      mock.patch.object(
+          database_session_service,
+          'async_sessionmaker',
+          side_effect=fake_async_sessionmaker,
+      ),
+  ):
+    database_session_service.DatabaseSessionService(
+        'postgresql+psycopg2://user:pass@localhost:5432/db'
+    )
+
+  assert captured_binds == [fake_engine, read_only_engine]
+  fake_engine.execution_options.assert_called_once_with(read_only=True)
+
+
 @pytest.mark.asyncio
 async def test_sqlite_session_service_accepts_sqlite_urls(
     tmp_path, monkeypatch
@@ -198,6 +267,67 @@ async def test_get_empty_session(session_service):
   )
 
 
+@pytest.mark.asyncio
+async def test_database_session_service_get_session_uses_read_only_factory():
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  service._prepare_tables = mock.AsyncMock()
+
+  read_only_session = mock.AsyncMock()
+  read_only_session.get = mock.AsyncMock(return_value=None)
+
+  @asynccontextmanager
+  async def fake_read_only_session():
+    yield read_only_session
+
+  service.database_session_factory = mock.Mock(
+      side_effect=AssertionError('write session factory should not be used')
+  )
+  service._read_only_database_session_factory = mock.Mock(
+      return_value=fake_read_only_session()
+  )
+
+  session = await service.get_session(
+      app_name='my_app', user_id='test_user', session_id='123'
+  )
+
+  assert session is None
+  service._read_only_database_session_factory.assert_called_once_with()
+  service.database_session_factory.assert_not_called()
+
+  await service.close()
+
+
+@pytest.mark.asyncio
+async def test_database_session_service_list_sessions_uses_read_only_factory():
+  service = DatabaseSessionService('sqlite+aiosqlite:///:memory:')
+  service._prepare_tables = mock.AsyncMock()
+
+  read_only_session = mock.AsyncMock()
+  empty_result = mock.Mock()
+  empty_result.scalars.return_value.all.return_value = []
+  read_only_session.execute = mock.AsyncMock(return_value=empty_result)
+  read_only_session.get = mock.AsyncMock(return_value=None)
+
+  @asynccontextmanager
+  async def fake_read_only_session():
+    yield read_only_session
+
+  service.database_session_factory = mock.Mock(
+      side_effect=AssertionError('write session factory should not be used')
+  )
+  service._read_only_database_session_factory = mock.Mock(
+      return_value=fake_read_only_session()
+  )
+
+  response = await service.list_sessions(app_name='my_app', user_id='test_user')
+
+  assert response.sessions == []
+  service._read_only_database_session_factory.assert_called_once_with()
+  service.database_session_factory.assert_not_called()
+
+  await service.close()
+
+
 @pytest.mark.asyncio
 async def test_create_get_session(session_service):
   app_name = 'my_app'