diff --git a/skyvern/config.py b/skyvern/config.py index aed4b70b..0646a5b6 100644 --- a/skyvern/config.py +++ b/skyvern/config.py @@ -26,6 +26,7 @@ class Settings(BaseSettings): MAX_RETRIES_PER_STEP: int = 5 DEBUG_MODE: bool = False DATABASE_STRING: str = "postgresql+psycopg://skyvern@localhost/skyvern" + DATABASE_STATEMENT_TIMEOUT_MS: int = 60000 PROMPT_ACTION_HISTORY_WINDOW: int = 1 TASK_RESPONSE_ACTION_SCREENSHOT_COUNT: int = 3 diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 76596713..a946a92d 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -74,12 +74,23 @@ from skyvern.webeye.actions.models import AgentStepOutput LOG = structlog.get_logger() +DB_CONNECT_ARGS: dict[str, Any] = {} + +if "postgresql+psycopg" in settings.DATABASE_STRING: + DB_CONNECT_ARGS = {"options": f"-c statement_timeout={settings.DATABASE_STATEMENT_TIMEOUT_MS}"} +elif "postgresql+asyncpg" in settings.DATABASE_STRING: + DB_CONNECT_ARGS = {"server_settings": {"statement_timeout": str(settings.DATABASE_STATEMENT_TIMEOUT_MS)}} + class AgentDB: def __init__(self, database_string: str, debug_enabled: bool = False) -> None: super().__init__() self.debug_enabled = debug_enabled - self.engine = create_async_engine(database_string, json_serializer=_custom_json_serializer, pool_pre_ping=True) + self.engine = create_async_engine( + database_string, + json_serializer=_custom_json_serializer, + connect_args=DB_CONNECT_ARGS, + ) self.Session = async_sessionmaker(bind=self.engine) async def create_task(