Merge pull request #1177 from PrefectHQ/database-timeouts

Add a db connection timeout; lengthen the command timeout

src/prefect/orion/database/configurations.py

@@ -8,7 +8,11 @@
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from typing_extensions import Literal
 
-from prefect.settings import PREFECT_ORION_DATABASE_ECHO, PREFECT_ORION_DATABASE_TIMEOUT
+from prefect.settings import (
+    PREFECT_ORION_DATABASE_CONNECTION_TIMEOUT,
+    PREFECT_ORION_DATABASE_ECHO,
+    PREFECT_ORION_DATABASE_TIMEOUT,
+)
 from prefect.utilities.asyncio import add_event_loop_shutdown_callback
 
 
@@ -25,10 +29,14 @@ def __init__(
         connection_url: str,
         echo: bool = None,
         timeout: float = None,
+        connection_timeout: float = None,
     ):
         self.connection_url = connection_url
         self.echo = echo or PREFECT_ORION_DATABASE_ECHO.value()
         self.timeout = timeout or PREFECT_ORION_DATABASE_TIMEOUT.value()
+        self.connection_timeout = (
+            connection_timeout or PREFECT_ORION_DATABASE_CONNECTION_TIMEOUT.value()
+        )
 
     def _unique_key(self) -> Tuple[Hashable, ...]:
         """
@@ -90,8 +98,16 @@ async def engine(self) -> sa.engine.Engine:
 
             # apply database timeout
             kwargs = dict()
+            connect_args = dict()
+
             if self.timeout is not None:
-                kwargs["connect_args"] = dict(command_timeout=self.timeout)
+                connect_args["command_timeout"] = self.timeout
+
+            if self.connection_timeout is not None:
+                connect_args["timeout"] = self.connection_timeout
+
+            if connect_args:
+                kwargs["connect_args"] = connect_args
 
             engine = create_async_engine(self.connection_url, echo=self.echo, **kwargs)
 

src/prefect/settings.py

@@ -284,9 +284,16 @@ def max_log_size_smaller_than_batch_size(values):
 
 PREFECT_ORION_DATABASE_TIMEOUT = Setting(
     Optional[float],
-    default=1,
+    default=10,
     description="""A statement timeout, in seconds, applied to all database
-    interactions made by the API. Defaults to `1`.""",
+    interactions made by the API. Defaults to `10`.""",
+)
+
+PREFECT_ORION_DATABASE_CONNECTION_TIMEOUT = Setting(
+    Optional[float],
+    default=5,
+    description="""A connection timeout, in seconds, applied to database
+    connections. Defaults to `5`.""",
 )
 
 PREFECT_ORION_SERVICES_RUN_IN_APP = Setting(