You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
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:
committed by
Copybara-Service
parent
125bc85ac5
commit
da73e718ef
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user