From bef1bcb2a4da24bd52a2e600d8c2929f2e15b02d Mon Sep 17 00:00:00 2001
From: Ruben Fiszel <ruben@rubenfiszel.com>
Date: Thu, 12 Sep 2024 18:49:43 +0200
Subject: [PATCH] fix: timeout in pg executor on postgresql connection after
 20s

---
 backend/windmill-worker/src/pg_executor.rs | 29 ++++++++++++++++------
 1 file changed, 21 insertions(+), 8 deletions(-)

diff --git a/backend/windmill-worker/src/pg_executor.rs b/backend/windmill-worker/src/pg_executor.rs
index f679d6c003..4f79c3a972 100644
--- a/backend/windmill-worker/src/pg_executor.rs
+++ b/backend/windmill-worker/src/pg_executor.rs
@@ -232,15 +232,21 @@ pub async fn do_postgresql(
                 .danger_accept_invalid_hostnames(true);
         }
 
-        let (client, connection) = tokio_postgres::connect(
-            &database_string,
-            MakeTlsConnector::new(connector.build().map_err(to_anyhow)?),
+        let (client, connection) = tokio::time::timeout(
+            std::time::Duration::from_secs(20),
+            tokio_postgres::connect(
+                &database_string,
+                MakeTlsConnector::new(connector.build().map_err(to_anyhow)?),
+            ),
         )
         .await
+        .map_err(to_anyhow)?
         .map_err(to_anyhow)?;
 
         let handle = tokio::spawn(async move {
-            if let Err(e) = connection.await {
+            if let Err(e) =
+                tokio::time::timeout(std::time::Duration::from_secs(20), connection).await
+            {
                 let mut mtex = CONNECTION_CACHE.lock().await;
                 *mtex = None;
                 tracing::error!("connection error: {}", e);
@@ -249,11 +255,18 @@ pub async fn do_postgresql(
         Some((client, handle))
     } else {
         tracing::info!("Creating new connection");
-        let (client, connection) = tokio_postgres::connect(&database_string, NoTls)
-            .await
-            .map_err(to_anyhow)?;
+        let (client, connection) = tokio::time::timeout(
+            std::time::Duration::from_secs(20),
+            tokio_postgres::connect(&database_string, NoTls),
+        )
+        .await
+        .map_err(to_anyhow)?
+        .map_err(to_anyhow)?;
+
         let handle = tokio::spawn(async move {
-            if let Err(e) = connection.await {
+            if let Err(e) =
+                tokio::time::timeout(std::time::Duration::from_secs(20), connection).await
+            {
                 let mut mtex = CONNECTION_CACHE.lock().await;
                 *mtex = None;
                 tracing::error!("connection error: {}", e);