Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions src/google/adk/sessions/database_session_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,6 +213,26 @@ async def _rollback_on_exception_session(
await sql_session.rollback()
raise

@asynccontextmanager
async def _readonly_session(
self,
) -> AsyncIterator[DatabaseSessionFactory]:
"""Yields a read-only database session for pure SELECT operations.

On Spanner this avoids OCC read-write transactions (which can trigger
RetryAborted when conflicting with concurrent writes). On PostgreSQL
the session is marked read-only. On other dialects the behaviour is
equivalent to a regular session but the transaction is never committed.
"""
async with self.database_session_factory() as sql_session:
conn = await sql_session.connection()
await conn.execution_options(postgresql_readonly=True)
try:
yield sql_session
except BaseException:
await sql_session.rollback()
raise

def _supports_row_level_locking(self) -> bool:
return self.db_engine.dialect.name in (
_MARIADB_DIALECT,
Expand Down Expand Up @@ -403,7 +423,7 @@ 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._readonly_session() as sql_session:
storage_session = await sql_session.get(
schema.StorageSession, (app_name, user_id, session_id)
)
Expand Down Expand Up @@ -458,7 +478,7 @@ 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._readonly_session() as sql_session:
stmt = select(schema.StorageSession).filter(
schema.StorageSession.app_name == app_name
)
Expand Down