Skip to content

Commit da73e71

Browse files
GWealecopybara-github
authored andcommitted
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
1 parent 125bc85 commit da73e71

File tree

2 files changed

+49
-0
lines changed

2 files changed

+49
-0
lines changed

src/google/adk/sessions/database_session_service.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,8 @@ def __init__(self, db_url: str, **kwargs: Any):
113113
connect_args = dict(engine_kwargs.get("connect_args", {}))
114114
connect_args.setdefault("check_same_thread", False)
115115
engine_kwargs["connect_args"] = connect_args
116+
elif url.get_backend_name() != "sqlite":
117+
engine_kwargs.setdefault("pool_pre_ping", True)
116118

117119
db_engine = create_async_engine(db_url, **engine_kwargs)
118120
if db_engine.dialect.name == "sqlite":

tests/unittests/sessions/test_session_service.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,12 @@
1616
from datetime import timezone
1717
import enum
1818
import sqlite3
19+
from unittest import mock
1920

2021
from google.adk.errors.already_exists_error import AlreadyExistsError
2122
from google.adk.events.event import Event
2223
from google.adk.events.event_actions import EventActions
24+
from google.adk.sessions import database_session_service
2325
from google.adk.sessions.base_session_service import GetSessionConfig
2426
from google.adk.sessions.database_session_service import DatabaseSessionService
2527
from google.adk.sessions.in_memory_session_service import InMemorySessionService
@@ -61,6 +63,51 @@ async def session_service(request, tmp_path):
6163
await service.close()
6264

6365

66+
def test_database_session_service_enables_pool_pre_ping_by_default():
67+
captured_kwargs = {}
68+
69+
def fake_create_async_engine(_db_url: str, **kwargs):
70+
captured_kwargs.update(kwargs)
71+
fake_engine = mock.Mock()
72+
fake_engine.dialect.name = 'postgresql'
73+
fake_engine.sync_engine = mock.Mock()
74+
return fake_engine
75+
76+
with mock.patch.object(
77+
database_session_service,
78+
'create_async_engine',
79+
side_effect=fake_create_async_engine,
80+
):
81+
database_session_service.DatabaseSessionService(
82+
'postgresql+psycopg2://user:pass@localhost:5432/db'
83+
)
84+
85+
assert captured_kwargs.get('pool_pre_ping') is True
86+
87+
88+
def test_database_session_service_respects_pool_pre_ping_override():
89+
captured_kwargs = {}
90+
91+
def fake_create_async_engine(_db_url: str, **kwargs):
92+
captured_kwargs.update(kwargs)
93+
fake_engine = mock.Mock()
94+
fake_engine.dialect.name = 'postgresql'
95+
fake_engine.sync_engine = mock.Mock()
96+
return fake_engine
97+
98+
with mock.patch.object(
99+
database_session_service,
100+
'create_async_engine',
101+
side_effect=fake_create_async_engine,
102+
):
103+
database_session_service.DatabaseSessionService(
104+
'postgresql+psycopg2://user:pass@localhost:5432/db',
105+
pool_pre_ping=False,
106+
)
107+
108+
assert captured_kwargs.get('pool_pre_ping') is False
109+
110+
64111
@pytest.mark.asyncio
65112
async def test_sqlite_session_service_accepts_sqlite_urls(
66113
tmp_path, monkeypatch

0 commit comments

Comments
 (0)