fix: Enable pool_pre_ping by default for non-SQLite database engines

This change sets `pool_pre_ping=True` in SQLAlchemy engine kwargs for database backends other than SQLite. This helps ensure that connections from the pool are still valid before being used, preventing issues with stale or disconnected connections. Tests are added to verify the default behavior and that explicit overrides are respected

Close #4211

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 864886767
This commit is contained in:
George Weale
2026-02-03 08:16:27 -08:00
committed by Copybara-Service
parent 125bc85ac5
commit da73e718ef
2 changed files with 49 additions and 0 deletions
@@ -113,6 +113,8 @@ class DatabaseSessionService(BaseSessionService):
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":
engine_kwargs.setdefault("pool_pre_ping", True)
db_engine = create_async_engine(db_url, **engine_kwargs)
if db_engine.dialect.name == "sqlite":
@@ -16,10 +16,12 @@ from datetime import datetime
from datetime import timezone
import enum
import sqlite3
from unittest import mock
from google.adk.errors.already_exists_error import AlreadyExistsError
from google.adk.events.event import Event
from google.adk.events.event_actions import EventActions
from google.adk.sessions import database_session_service
from google.adk.sessions.base_session_service import GetSessionConfig
from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -61,6 +63,51 @@ async def session_service(request, tmp_path):
await service.close()
def test_database_session_service_enables_pool_pre_ping_by_default():
captured_kwargs = {}
def fake_create_async_engine(_db_url: str, **kwargs):
captured_kwargs.update(kwargs)
fake_engine = mock.Mock()
fake_engine.dialect.name = 'postgresql'
fake_engine.sync_engine = mock.Mock()
return fake_engine
with mock.patch.object(
database_session_service,
'create_async_engine',
side_effect=fake_create_async_engine,
):
database_session_service.DatabaseSessionService(
'postgresql+psycopg2://user:pass@localhost:5432/db'
)
assert captured_kwargs.get('pool_pre_ping') is True
def test_database_session_service_respects_pool_pre_ping_override():
captured_kwargs = {}
def fake_create_async_engine(_db_url: str, **kwargs):
captured_kwargs.update(kwargs)
fake_engine = mock.Mock()
fake_engine.dialect.name = 'postgresql'
fake_engine.sync_engine = mock.Mock()
return fake_engine
with mock.patch.object(
database_session_service,
'create_async_engine',
side_effect=fake_create_async_engine,
):
database_session_service.DatabaseSessionService(
'postgresql+psycopg2://user:pass@localhost:5432/db',
pool_pre_ping=False,
)
assert captured_kwargs.get('pool_pre_ping') is False
@pytest.mark.asyncio
async def test_sqlite_session_service_accepts_sqlite_urls(
tmp_path, monkeypatch