vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix statement timeout issue (#4291)

Browse Source
This commit is contained in:
Shuchang Zheng
2025-12-12 22:39:33 -08:00
committed by GitHub
parent 3d54a288ad
commit 6fffa5cc02
1 changed files with 10 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
skyvern/forge/sdk/db/agent_db.py
View File

@@ -175,8 +175,10 @@ def _connect_args_for_driver(database_string: str) -> dict[str, Any]:
if settings.DB_DISABLE_PREPARED_STATEMENTS: if settings.DB_DISABLE_PREPARED_STATEMENTS:
if driver == "postgresql+psycopg": if driver == "postgresql+psycopg":
# psycopg3: disable server-side prepares # psycopg3: completely disable prepared statements (None vs 0)
args["prepare_threshold"] = 0 # 0 disables caching but still uses prepared statements
# None completely disables prepared statement usage
args["prepare_threshold"] = None
elif driver == "postgresql+asyncpg": elif driver == "postgresql+asyncpg":
# asyncpg: disable statement cache (prepared statements) # asyncpg: disable statement cache (prepared statements)
args["statement_cache_size"] = 0 args["statement_cache_size"] = 0
@@ -192,11 +194,16 @@ def _install_statement_timeout(engine: AsyncEngine, timeout_ms: int) -> None:
if not timeout_ms or timeout_ms <= 0: if not timeout_ms or timeout_ms <= 0:
return return
timeout_value = int(timeout_ms)
# Works for direct AND poolers because it's not a startup parameter. # Works for direct AND poolers because it's not a startup parameter.
# Applies per-transaction, which is the most reliable behavior with transaction pooling. # Applies per-transaction, which is the most reliable behavior with transaction pooling.
@event.listens_for(engine.sync_engine, "begin") @event.listens_for(engine.sync_engine, "begin")
def _set_timeout(conn: Connection) -> None: def _set_timeout(conn: Connection) -> None:
conn.exec_driver_sql(f"SET LOCAL statement_timeout = {int(timeout_ms)}") # Use a unique comment with object id to prevent psycopg3 from reusing prepared statement names,
# which can cause "prepared statement already exists" errors with connection poolers.
sql = f"SET LOCAL statement_timeout = {timeout_value} /* {id(conn)} */"
conn.exec_driver_sql(sql)
def make_async_engine(database_string: str) -> AsyncEngine: def make_async_engine(database_string: str) -> AsyncEngine: