introduce db timeout (#1222)

Skyvern-AI · Nov 20, 2024 · e4c1cf6 · e4c1cf6
1 parent 9896a70
commit e4c1cf6
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git 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
@@ -74,12 +74,23 @@
 
 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(