diff --git a/skyvern/config.py b/skyvern/config.py index 93f3506ea..1825e580f 100644 --- a/skyvern/config.py +++ b/skyvern/config.py @@ -155,6 +155,11 @@ class Settings(BaseSettings): DISABLE_CONNECTION_POOL: bool = False DATABASE_POOL_SIZE: int = 5 DATABASE_POOL_MAX_OVERFLOW: int = 10 + # Timeout/recycle defaults mirror SQLAlchemy's QueuePool. Size pools per service + # via env vars: raising defaults here multiplies across every engine and replica + # and can exhaust pgbouncer client connections. + DATABASE_POOL_TIMEOUT: int = 30 + DATABASE_POOL_RECYCLE: int = -1 PROMPT_ACTION_HISTORY_WINDOW: int = 1 TASK_RESPONSE_ACTION_SCREENSHOT_COUNT: int = 3 diff --git a/skyvern/forge/sdk/db/agent_db.py b/skyvern/forge/sdk/db/agent_db.py index a382049e7..d5d1dde3a 100644 --- a/skyvern/forge/sdk/db/agent_db.py +++ b/skyvern/forge/sdk/db/agent_db.py @@ -348,6 +348,8 @@ def _build_engine(database_string: str) -> AsyncEngine: pool_pre_ping=True, pool_size=settings.DATABASE_POOL_SIZE, max_overflow=settings.DATABASE_POOL_MAX_OVERFLOW, + pool_timeout=settings.DATABASE_POOL_TIMEOUT, + pool_recycle=settings.DATABASE_POOL_RECYCLE, ) diff --git a/tests/unit/test_db_engine_connect_args.py b/tests/unit/test_db_engine_connect_args.py index 02f2af3ae..3cb466352 100644 --- a/tests/unit/test_db_engine_connect_args.py +++ b/tests/unit/test_db_engine_connect_args.py @@ -197,6 +197,47 @@ def test_build_engine_keeps_pool_class_selection_independent_of_connect_args(mon assert disabled_kwargs["connect_args"]["prepared_statement_name_func"] is agent_db._asyncpg_prepared_statement_name +def test_build_engine_forwards_pool_sizing_settings(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[dict[str, Any]] = [] + + def fake_create_async_engine(database_string: str, **kwargs: Any) -> object: + calls.append({"database_string": database_string, "kwargs": kwargs}) + return object() + + monkeypatch.setattr(agent_db, "create_async_engine", fake_create_async_engine) + monkeypatch.setattr(settings, "DISABLE_CONNECTION_POOL", False) + monkeypatch.setattr(settings, "DATABASE_POOL_SIZE", 20) + monkeypatch.setattr(settings, "DATABASE_POOL_MAX_OVERFLOW", 20) + monkeypatch.setattr(settings, "DATABASE_POOL_TIMEOUT", 10) + monkeypatch.setattr(settings, "DATABASE_POOL_RECYCLE", 1800) + + agent_db._build_engine(DIRECT_ASYNCPG) + + kwargs = calls[-1]["kwargs"] + assert kwargs["pool_size"] == 20 + assert kwargs["max_overflow"] == 20 + assert kwargs["pool_timeout"] == 10 + assert kwargs["pool_recycle"] == 1800 + + +def test_build_engine_null_pool_does_not_forward_queue_pool_params(monkeypatch: pytest.MonkeyPatch) -> None: + calls: list[dict[str, Any]] = [] + + def fake_create_async_engine(database_string: str, **kwargs: Any) -> object: + calls.append({"database_string": database_string, "kwargs": kwargs}) + return object() + + monkeypatch.setattr(agent_db, "create_async_engine", fake_create_async_engine) + monkeypatch.setattr(settings, "DISABLE_CONNECTION_POOL", True) + + agent_db._build_engine(DIRECT_ASYNCPG) + + kwargs = calls[-1]["kwargs"] + assert kwargs["poolclass"] is pool.NullPool + for queue_pool_param in ("pool_size", "max_overflow", "pool_timeout", "pool_recycle"): + assert queue_pool_param not in kwargs + + def test_build_engine_strips_asyncpg_weak_sslmode_when_requiring_pooler_ssl(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[dict[str, Any]] = []