Merge branch 'main' into restricted-job-tokens

backend/tests/worker.rs

@@ -1731,6 +1731,78 @@ async fn test_postgresql_cached_connection_resets_session(
     Ok(())
 }
 
+/// The idle cached connection must not hold the server slot a job's own
+/// connection needs. A role limited to one connection stands in for a
+/// session-mode pooler with one slot, where the second job would wait forever
+/// instead of failing.
+#[sqlx::test(fixtures("base"))]
+#[serial(pg_cache)]
+async fn test_postgresql_cached_connection_released_for_other_key(
+    db: Pool<Postgres>,
+) -> anyhow::Result<()> {
+    use windmill_worker::pg_executor::clear_pg_cache;
+
+    initialize_tracing().await;
+    clear_pg_cache().await;
+
+    sqlx::query(
+        "DO $$ BEGIN
+           IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = 'wm_pg_cache_one_conn') THEN
+             CREATE ROLE wm_pg_cache_one_conn LOGIN PASSWORD 'changeme' CONNECTION LIMIT 1;
+           END IF;
+         END $$",
+    )
+    .execute(&db)
+    .await?;
+
+    let server = ApiServer::start(db.clone()).await?;
+    let port = server.addr.port();
+
+    // Two cache keys for the same role: only the sslmode differs.
+    let run = |sslmode: &str| {
+        RunJob::from(JobPayload::Code(RawCode {
+            hash: None,
+            content: "SELECT 1 as n;".into(),
+            path: None,
+            lock: None,
+            language: ScriptLang::Postgresql,
+            cache_ttl: None,
+            cache_ignore_s3_path: None,
+            dedicated_worker: None,
+            concurrency_settings: windmill_common::runnable_settings::ConcurrencySettings::default(
+            )
+            .into(),
+            debouncing_settings: windmill_common::runnable_settings::DebouncingSettings::default(),
+            modules: None,
+            tag: None,
+        }))
+        .arg(
+            "database",
+            json!({"host": "localhost", "port": 5432, "dbname": "windmill",
+                   "user": "wm_pg_cache_one_conn", "password": "changeme", "sslmode": sslmode}),
+        )
+        .run_until_complete(&db, false, port)
+    };
+
+    // The evicted connection's backend exits asynchronously, so a fresh
+    // connection can briefly still count it. Retrying absorbs that; without the
+    // eviction the cached connection stays open for 60s and every retry fails.
+    for sslmode in ["disable", "prefer", "disable"] {
+        let mut result = json!(null);
+        for _ in 0..5 {
+            result = run(sslmode).await.json_result().unwrap();
+            if result == json!([{"n": 1}]) {
+                break;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(200)).await;
+        }
+        assert_eq!(result, json!([{"n": 1}]), "sslmode={sslmode}");
+    }
+
+    clear_pg_cache().await;
+    Ok(())
+}
+
 /// Runs multiple PG jobs through a SINGLE worker (like production) to verify
 /// that SET ROLE / search_path changes do not leak across jobs.
 #[sqlx::test(fixtures("base"))]

backend/windmill-common/Cargo.toml

@@ -22,7 +22,7 @@ smtp = ["dep:mail-send", "dep:rustls", "dep:tokio-rustls", "dep:rustls-native-ce
 scoped_cache = []
 cloud = []
 dev_override = []
-openidconnect = ["dep:openidconnect"]
+openidconnect = ["dep:openidconnect", "dep:rsa"]
 python = ["dep:windmill-parser-py"]
 
 [lib]

backend/windmill-common/src/lib.rs

@@ -1069,6 +1069,18 @@ impl TokioPgConnection {
     }
 }
 
+/// Without these, a server that vanishes without closing the socket (a failover,
+/// a dropped route) leaves a query waiting on a read for the OS default of two
+/// hours. The server's kernel answers the probes, so a slow query is unaffected.
+fn set_pg_keepalive(config: &mut tokio_postgres::Config) {
+    config
+        .keepalives(true)
+        .keepalives_idle(std::time::Duration::from_secs(60))
+        .keepalives_interval(std::time::Duration::from_secs(10))
+        .keepalives_retries(6)
+        .tcp_user_timeout(std::time::Duration::from_secs(120));
+}
+
 impl PgDatabase {
     /// The role the connection logs in as, whichever way it authenticates.
     pub fn login_name(&self) -> &str {
@@ -1110,6 +1122,12 @@ impl PgDatabase {
         self.options.as_deref().filter(|o| !o.is_empty())
     }
 
+    fn uri_config(&self) -> Result<tokio_postgres::Config, error::Error> {
+        let mut config: tokio_postgres::Config = self.to_uri().parse().map_err(to_anyhow)?;
+        set_pg_keepalive(&mut config);
+        Ok(config)
+    }
+
     pub async fn connect(
         &self,
         main_db: Option<&DB>,
@@ -1256,10 +1274,8 @@ impl PgDatabase {
 
             let (client, connection) = tokio::time::timeout(
                 std::time::Duration::from_secs(20),
-                tokio_postgres::connect(
-                    &self.to_uri(),
-                    MakeTlsConnector::new(connector.build().map_err(to_anyhow)?),
-                ),
+                self.uri_config()?
+                    .connect(MakeTlsConnector::new(connector.build().map_err(to_anyhow)?)),
             )
             .await
             .map_err(to_anyhow)?
@@ -1270,7 +1286,7 @@ impl PgDatabase {
             tracing::info!("Creating new connection");
             let (client, connection) = tokio::time::timeout(
                 std::time::Duration::from_secs(20),
-                tokio_postgres::connect(&self.to_uri(), NoTls),
+                self.uri_config()?.connect(NoTls),
             )
             .await
             .map_err(to_anyhow)?
@@ -1387,6 +1403,7 @@ impl PgDatabase {
         if let Some(options) = self.non_empty_options() {
             config.options(options);
         }
+        set_pg_keepalive(&mut config);
 
         let (client, connection) = tokio::time::timeout(
             std::time::Duration::from_secs(20),

backend/windmill-worker/src/pg_executor.rs

@@ -13,7 +13,7 @@ use rust_decimal::{prelude::FromPrimitive, Decimal};
 use serde_json::value::RawValue;
 use serde_json::Map;
 use serde_json::Value;
-use tokio::sync::{Mutex, RwLock};
+use tokio::sync::RwLock;
 use tokio_postgres::Client;
 use tokio_postgres::{types::ToSql, Row};
 use tokio_postgres::{
@@ -54,16 +54,266 @@ use windmill_common::client::AuthedClient;
 use windmill_types::s3::S3Object;
 
 lazy_static! {
-    pub static ref CONNECTION_CACHE: Arc<Mutex<Option<(String, tokio_postgres::Client)>>> =
-        Arc::new(Mutex::new(None));
     pub static ref CONNECTION_COUNTER: Arc<RwLock<HashMap<String, u64>>> =
         Arc::new(RwLock::new(HashMap::new()));
-    pub static ref LAST_QUERY: AtomicU64 = AtomicU64::new(0);
     pub static ref CACHE_HITS: AtomicU64 = AtomicU64::new(0);
 }
 
+/// How far a job's statements have got, read by its stall warning.
+struct PgProgress {
+    start: std::time::Instant,
+    statement: AtomicUsize,
+    rows: AtomicU64,
+    last_progress_ms: AtomicU64,
+}
+
+impl PgProgress {
+    fn new() -> Self {
+        Self {
+            start: std::time::Instant::now(),
+            statement: AtomicUsize::new(0),
+            rows: AtomicU64::new(0),
+            last_progress_ms: AtomicU64::new(0),
+        }
+    }
+
+    fn elapsed_ms(&self) -> u64 {
+        self.start.elapsed().as_millis() as u64
+    }
+
+    fn start_statement(&self, index: usize) {
+        self.statement.store(index, Ordering::Relaxed);
+        self.last_progress_ms
+            .store(self.elapsed_ms(), Ordering::Relaxed);
+    }
+
+    fn row(&self) {
+        self.rows.fetch_add(1, Ordering::Relaxed);
+        self.last_progress_ms
+            .store(self.elapsed_ms(), Ordering::Relaxed);
+    }
+
+    fn rows(&self) -> u64 {
+        self.rows.load(Ordering::Relaxed)
+    }
+}
+
+const PG_STALL_WARNING_AFTER: Duration = Duration::from_secs(5 * 60);
+/// Steps faster than this are not logged: a log line is a write to the main DB,
+/// and a script split into many statements would pay one per statement.
+const PG_SLOW_STEP: Duration = Duration::from_secs(1);
+
+fn connection_kind(fresh_connection: bool) -> &'static str {
+    if fresh_connection {
+        "new connection"
+    } else {
+        "worker's cached connection"
+    }
+}
+
+/// Warns once per stall in the job log; never completes. A statement still
+/// computing its first row also trips it, which is why it only warns: the job
+/// timeout stays the one thing that stops a job.
+async fn warn_on_stalled_statement(
+    progress: &PgProgress,
+    statement_count: usize,
+    fresh_connection: bool,
+    job_id: Uuid,
+    workspace_id: &str,
+    conn: &Connection,
+) -> std::convert::Infallible {
+    let mut warned_for = None;
+    loop {
+        tokio::time::sleep(Duration::from_secs(30)).await;
+        let last = progress.last_progress_ms.load(Ordering::Relaxed);
+        let stalled = Duration::from_millis(progress.elapsed_ms().saturating_sub(last));
+        if stalled < PG_STALL_WARNING_AFTER || warned_for == Some(last) {
+            continue;
+        }
+        warned_for = Some(last);
+        windmill_queue::append_logs(
+            &job_id,
+            workspace_id,
+            format!(
+                "No new row for {} min on statement {}/{statement_count} ({} rows so far, {}). \
+                 A query still running on the database is normal; if the database shows no \
+                 active query for this connection, the job is stuck in the worker.\n",
+                stalled.as_secs() / 60,
+                progress.statement.load(Ordering::Relaxed) + 1,
+                progress.rows(),
+                connection_kind(fresh_connection),
+            ),
+            conn,
+        )
+        .await;
+    }
+}
+
+/// One reusable connection per worker process, checked out by a job for the
+/// duration of its query and checked back in afterwards.
+///
+/// Invariant: the process never keeps an idle connection while one of its jobs
+/// runs on a connection of its own. Behind a session-mode pooler (PgBouncer,
+/// RDS Proxy pinning, …) the idle connection holds a server slot, and the job
+/// on its own connection waits for that slot with no timeout, forever if the
+/// idle connection keeps being reused by other jobs.
+static PG_CONNECTION_CACHE: std::sync::Mutex<PgConnectionCache> =
+    std::sync::Mutex::new(PgConnectionCache { idle: None, checked_out: false, uncached: 0 });
+static PG_CONNECTION_CACHE_SWEEPER: std::sync::Once = std::sync::Once::new();
+const PG_CONNECTION_IDLE_TIMEOUT: Duration = Duration::from_secs(60);
+
+struct PgConnectionCache {
+    idle: Option<PgConnection>,
+    /// The cached connection is out with a running job.
+    checked_out: bool,
+    /// Jobs of this process running on a connection that is not the cached one.
+    uncached: usize,
+}
+
+fn pg_connection_cache() -> std::sync::MutexGuard<'static, PgConnectionCache> {
+    PG_CONNECTION_CACHE
+        .lock()
+        .unwrap_or_else(|e| e.into_inner())
+}
+
+struct PgConnection {
+    key: String,
+    client: Client,
+    connection_task: tokio::task::JoinHandle<()>,
+    last_used: std::time::Instant,
+}
+
+impl Drop for PgConnection {
+    fn drop(&mut self) {
+        self.connection_task.abort();
+    }
+}
+
+#[derive(Clone, Copy, PartialEq)]
+enum PgLeaseSlot {
+    CheckedOut,
+    Uncached,
+}
+
+/// A job's claim on a connection. Dropping it checks the connection back in,
+/// or closes it when keeping it would break the cache invariant.
+struct PgConnectionLease {
+    slot: PgLeaseSlot,
+    conn: Option<PgConnection>,
+    /// Cache this job's own connection on release.
+    cache_on_release: bool,
+}
+
+impl PgConnectionLease {
+    /// Takes the cached connection when it matches `key`. Otherwise registers
+    /// the job as uncached and closes the idle connection, whatever its key, so
+    /// it cannot hold the server slot the job's own connection is about to need.
+    fn checkout(key: &str) -> Self {
+        let (lease, evicted) = {
+            let mut cache = pg_connection_cache();
+            match cache.idle.take() {
+                Some(conn) if conn.key == key => {
+                    cache.checked_out = true;
+                    (Self::new(PgLeaseSlot::CheckedOut, Some(conn)), None)
+                }
+                other => {
+                    cache.uncached += 1;
+                    (Self::new(PgLeaseSlot::Uncached, None), other)
+                }
+            }
+        };
+        drop(evicted);
+        lease
+    }
+
+    fn uncached() -> Self {
+        pg_connection_cache().uncached += 1;
+        Self::new(PgLeaseSlot::Uncached, None)
+    }
+
+    fn new(slot: PgLeaseSlot, conn: Option<PgConnection>) -> Self {
+        Self { slot, conn, cache_on_release: false }
+    }
+
+    /// Gives up a cached connection that failed its probe; the job then
+    /// connects on its own.
+    fn discard(&mut self) {
+        let conn = self.conn.take();
+        if self.slot == PgLeaseSlot::CheckedOut {
+            let mut cache = pg_connection_cache();
+            cache.checked_out = false;
+            cache.uncached += 1;
+            self.slot = PgLeaseSlot::Uncached;
+        }
+        drop(conn);
+    }
+
+    fn client(&self) -> &Client {
+        &self.conn.as_ref().expect("lease holds a connection").client
+    }
+}
+
+impl Drop for PgConnectionLease {
+    fn drop(&mut self) {
+        let Some(mut conn) = self.conn.take() else {
+            if self.slot == PgLeaseSlot::Uncached {
+                pg_connection_cache().uncached -= 1;
+            }
+            return;
+        };
+        let kept = {
+            let mut cache = pg_connection_cache();
+            let wants_cache = match self.slot {
+                PgLeaseSlot::CheckedOut => {
+                    cache.checked_out = false;
+                    true
+                }
+                PgLeaseSlot::Uncached => {
+                    cache.uncached -= 1;
+                    self.cache_on_release
+                }
+            };
+            if wants_cache && cache.uncached == 0 && !cache.checked_out && cache.idle.is_none() {
+                conn.last_used = std::time::Instant::now();
+                cache.idle = Some(conn);
+                None
+            } else {
+                Some(conn)
+            }
+        };
+        match kept {
+            None => PG_CONNECTION_CACHE_SWEEPER.call_once(|| {
+                tokio::spawn(sweep_idle_pg_connection());
+            }),
+            Some(conn) => drop(conn),
+        }
+    }
+}
+
+async fn sweep_idle_pg_connection() {
+    loop {
+        tokio::time::sleep(Duration::from_secs(5)).await;
+        let expired = {
+            let mut cache = pg_connection_cache();
+            if cache
+                .idle
+                .as_ref()
+                .is_some_and(|c| c.last_used.elapsed() > PG_CONNECTION_IDLE_TIMEOUT)
+            {
+                cache.idle.take()
+            } else {
+                None
+            }
+        };
+        if expired.is_some() {
+            tracing::info!("Closing cache pg executor connection due to inactivity");
+        }
+    }
+}
+
 pub async fn clear_pg_cache() {
-    *CONNECTION_CACHE.lock().await = None;
+    let idle = pg_connection_cache().idle.take();
+    drop(idle);
     CONNECTION_COUNTER.write().await.clear();
 }
 
@@ -158,8 +408,6 @@ async fn new_pg_connection(
     };
     let handle = tokio::spawn(async move {
         if let Err(e) = connection.await {
-            let mut mtex = CONNECTION_CACHE.lock().await;
-            *mtex = None;
             tracing::error!("connection error: {}", e);
         }
     });
@@ -483,6 +731,7 @@ fn do_postgresql_inner<'a>(
     workspace_id: &'a str,
     log_conn: &'a Connection,
     raw_output: bool,
+    progress: &'a PgProgress,
 ) -> error::Result<BoxFuture<'a, error::Result<Vec<Box<RawValue>>>>> {
     let mut query_params = vec![];
     let mut param_types: Vec<Type> = vec![];
@@ -633,10 +882,13 @@ fn do_postgresql_inner<'a>(
 
         if skip_collect {
             futures::pin_mut!(rows);
-            while rows.try_next().await.map_err(to_anyhow)?.is_some() {}
+            while rows.try_next().await.map_err(to_anyhow)?.is_some() {
+                progress.row();
+            }
         } else if let Some(ref s3) = s3 {
             let format_state_ref = &format_state;
             let rows_stream = rows.map_err(to_anyhow).map(move |row_result| {
+                progress.row();
                 row_result.and_then(|row| {
                     postgres_row_to_json_value_with_state(row, format_state_ref).map_err(to_anyhow)
                 })
@@ -674,6 +926,7 @@ fn do_postgresql_inner<'a>(
             let mut column_names: Option<Vec<String>> = None;
 
             while let Some(row) = rows.try_next().await.map_err(to_anyhow)? {
+                progress.row();
                 if column_names.is_none() {
                     column_names = Some(
                         row.columns()
@@ -868,97 +1121,107 @@ pub async fn do_postgresql(
         database.to_uri(),
         auth_mode.cache_key_segment()
     );
-    let database_string_clone = database_string.clone();
 
-    let cached_client;
-    let new_client;
-    if !*CLOUD_HOSTED {
-        let mut guard = CONNECTION_CACHE.try_lock().ok();
+    let connect_started = std::time::Instant::now();
+    let mut cached_connection_failed_reset = false;
+    let mut lease = if *CLOUD_HOSTED {
+        PgConnectionLease::uncached()
+    } else {
         increment_connection_counter(&database_string).await;
+        PgConnectionLease::checkout(&database_string)
+    };
 
-        if guard
-            .as_ref()
-            .is_some_and(|x| x.as_ref().is_some_and(|y| y.0 == database_string))
-        {
-            // Probe the cached connection with a curated session reset before
-            // reusing it. Each statement targets a specific class of state:
-            //
-            //   RESET ALL                     — GUC parameters (search_path,
-            //                                   application_name, statement_
-            //                                   timeout, transaction_*…). Note
-            //                                   that this does NOT reset SET
-            //                                   ROLE or SET SESSION
-            //                                   AUTHORIZATION (security!).
-            //   RESET SESSION AUTHORIZATION   — undoes both `SET SESSION
-            //                                   AUTHORIZATION` and `SET ROLE`,
-            //                                   restoring the connecting user.
-            //                                   Without this a previous job
-            //                                   leaving an elevated role
-            //                                   active would silently leak
-            //                                   permissions into the next.
-            //   UNLISTEN *                    — drops LISTEN registrations.
-            //   CLOSE ALL                     — closes open cursors.
-            //   pg_advisory_unlock_all()      — releases any session-scoped
-            //                                   advisory locks. Without this
-            //                                   a job that called
-            //                                   pg_advisory_lock and exited
-            //                                   without unlocking would block
-            //                                   later jobs holding the same
-            //                                   key (DISCARD ALL covered this
-            //                                   too).
-            //
-            // We deliberately do NOT use `DISCARD ALL`. DISCARD includes
-            // `DEALLOCATE ALL`, which deallocates *all* prepared statements
-            // server-side — including the typeinfo statements that
-            // tokio_postgres caches per-Client to resolve custom enum/domain
-            // Oids. After DISCARD, tokio_postgres still holds Statement
-            // objects whose names the server has forgotten, so the next
-            // custom-type query fails with `prepared statement "sN" does not
-            // exist`. The trade-off: temp tables and user-PREPARE statements
-            // may persist across cached-connection reuse (rare in datatable /
-            // script workloads).
-            //
-            // Doubles as a liveness probe — if the connection is broken any
-            // statement in the chain fails and we replace it.
-            let probe_client = &guard.as_ref().unwrap().as_ref().unwrap().1;
-            if probe_client
-                .batch_execute(
-                    "RESET ALL; \
+    if lease.conn.is_some() {
+        // Probe the cached connection with a curated session reset before
+        // reusing it. Each statement targets a specific class of state:
+        //
+        //   RESET ALL                     — GUC parameters (search_path,
+        //                                   application_name, statement_
+        //                                   timeout, transaction_*…). Note
+        //                                   that this does NOT reset SET
+        //                                   ROLE or SET SESSION
+        //                                   AUTHORIZATION (security!).
+        //   RESET SESSION AUTHORIZATION   — undoes both `SET SESSION
+        //                                   AUTHORIZATION` and `SET ROLE`,
+        //                                   restoring the connecting user.
+        //                                   Without this a previous job
+        //                                   leaving an elevated role
+        //                                   active would silently leak
+        //                                   permissions into the next.
+        //   UNLISTEN *                    — drops LISTEN registrations.
+        //   CLOSE ALL                     — closes open cursors.
+        //   pg_advisory_unlock_all()      — releases any session-scoped
+        //                                   advisory locks. Without this
+        //                                   a job that called
+        //                                   pg_advisory_lock and exited
+        //                                   without unlocking would block
+        //                                   later jobs holding the same
+        //                                   key (DISCARD ALL covered this
+        //                                   too).
+        //
+        // We deliberately do NOT use `DISCARD ALL`. DISCARD includes
+        // `DEALLOCATE ALL`, which deallocates *all* prepared statements
+        // server-side — including the typeinfo statements that
+        // tokio_postgres caches per-Client to resolve custom enum/domain
+        // Oids. After DISCARD, tokio_postgres still holds Statement
+        // objects whose names the server has forgotten, so the next
+        // custom-type query fails with `prepared statement "sN" does not
+        // exist`. The trade-off: temp tables and user-PREPARE statements
+        // may persist across cached-connection reuse (rare in datatable /
+        // script workloads).
+        //
+        // Doubles as a liveness probe — if the connection is broken any
+        // statement in the chain fails and we replace it.
+        if lease
+            .client()
+            .batch_execute(
+                "RESET ALL; \
                      RESET SESSION AUTHORIZATION; \
                      UNLISTEN *; \
                      CLOSE ALL; \
                      SELECT pg_advisory_unlock_all();",
-                )
-                .await
-                .is_ok()
-            {
-                tracing::info!("Using cached connection");
-                CACHE_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-                LAST_QUERY.store(
-                    chrono::Utc::now().timestamp().try_into().unwrap_or(0),
-                    std::sync::atomic::Ordering::Relaxed,
-                );
-                cached_client = guard;
-                new_client = None;
-            } else {
-                tracing::info!("Cached connection is stale, creating new one");
-                if let Some(ref mut g) = guard {
-                    **g = None;
-                }
-                drop(guard);
-                cached_client = None;
-                new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?);
-            }
+            )
+            .await
+            .is_ok()
+        {
+            tracing::info!("Using cached connection");
+            CACHE_HITS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
         } else {
-            // Release the lock before connecting so the post-query caching
-            // code can re-acquire it.
-            drop(guard);
-            cached_client = None;
-            new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?);
+            tracing::info!("Cached connection is stale, creating new one");
+            cached_connection_failed_reset = true;
+            lease.discard();
         }
-    } else {
-        cached_client = None;
-        new_client = Some(new_pg_connection(&database, auth_mode, conn.as_sql()).await?);
+    }
+    if lease.conn.is_none() {
+        let (client, connection_task) =
+            new_pg_connection(&database, auth_mode, conn.as_sql()).await?;
+        lease.conn = Some(PgConnection {
+            key: database_string.clone(),
+            client,
+            connection_task,
+            last_used: std::time::Instant::now(),
+        });
+    }
+
+    let fresh_connection = lease.slot == PgLeaseSlot::Uncached;
+    let log_progress = !run_inline && !annotations.prepare;
+    let connect_time = connect_started.elapsed();
+    if log_progress && connect_time >= PG_SLOW_STEP {
+        windmill_queue::append_logs(
+            &job.id,
+            &job.workspace_id,
+            format!(
+                "Getting a database connection took {} ms ({})\n",
+                connect_time.as_millis(),
+                if cached_connection_failed_reset {
+                    "the cached connection failed its reset, then a new connection"
+                } else {
+                    connection_kind(fresh_connection)
+                }
+            ),
+            conn,
+        )
+        .await;
     }
 
     let (mut sig, _) = parse_pgsql_sig_with_typed_schema(&query)
@@ -985,12 +1248,7 @@ pub async fn do_postgresql(
 
     let queries = parse_sql_blocks(query, true);
 
-    let (client, handle) = if let Some((client, handle)) = new_client.as_ref() {
-        (client, Some(handle))
-    } else {
-        let (_, client) = cached_client.as_ref().unwrap().as_ref().unwrap();
-        (client, None)
-    };
+    let client = lease.client();
 
     let param_idx_to_arg_and_value = sig
         .args
@@ -1000,6 +1258,9 @@ pub async fn do_postgresql(
 
     let size = AtomicUsize::new(0);
     let size_ref = &size;
+    let progress = PgProgress::new();
+    let progress_ref = &progress;
+    let statement_count = queries.len();
     let result_f = async move {
         let mut results = vec![];
         // Session reset (DISCARD ALL) is now handled eagerly when validating
@@ -1037,6 +1298,9 @@ pub async fn do_postgresql(
             let skip_collect = collection_strategy.collect_last_statement_only(queries.len())
                 && i < queries.len() - 1;
             let is_last = i == queries.len() - 1;
+            progress_ref.start_statement(i);
+            let rows_before = progress_ref.rows();
+            let statement_started = std::time::Instant::now();
             let result = do_postgresql_inner(
                 query.to_string(),
                 &param_idx_to_arg_and_value,
@@ -1058,8 +1322,24 @@ pub async fn do_postgresql(
                 &job.workspace_id,
                 conn,
                 annotations.raw_output && is_last && !skip_collect,
+                progress_ref,
             )?
             .await?;
+            let statement_time = statement_started.elapsed();
+            if log_progress && statement_time >= PG_SLOW_STEP {
+                windmill_queue::append_logs(
+                    &job.id,
+                    &job.workspace_id,
+                    format!(
+                        "Statement {}/{statement_count}: {} rows in {} ms\n",
+                        i + 1,
+                        progress_ref.rows() - rows_before,
+                        statement_time.as_millis()
+                    ),
+                    conn,
+                )
+                .await;
+            }
             results.push(result);
         }
 
@@ -1072,6 +1352,16 @@ pub async fn do_postgresql(
             collection_strategy.collect(results)
         }
     };
+    let result_f = async {
+        if log_progress {
+            tokio::select! {
+                result = result_f => result,
+                never = warn_on_stalled_statement(progress_ref, statement_count, fresh_connection, job.id, &job.workspace_id, conn) => match never {},
+            }
+        } else {
+            result_f.await
+        }
+    };
 
     let result = if run_inline {
         result_f.await
@@ -1092,80 +1382,13 @@ pub async fn do_postgresql(
     }
     .map_err(|e| map_s3object_jsonb_overflow(e, had_s3object_input))?;
 
-    // Release the cache lock now that we have the result — allows the
-    // post-query caching code below to re-acquire it if needed.
-    drop(cached_client);
-
     *mem_peak = size.load(Ordering::Relaxed) as i32;
 
-    if let Some(handle) = handle {
-        if !*CLOUD_HOSTED {
-            if let Ok(mut mtex) = CONNECTION_CACHE.try_lock() {
-                if mtex.as_ref().is_none_or(|x| x.0 != database_string) {
-                    let abort_handler = handle.abort_handle();
-
-                    let mut cache_new_con = false;
-                    if let Some(new_client) = new_client {
-                        cache_new_con = is_most_used_conn(&database_string).await;
-                        if cache_new_con {
-                            *mtex = Some((database_string, new_client.0));
-                        } else {
-                            new_client.1.abort();
-                        }
-                    } else {
-                        handle.abort();
-                    }
-
-                    if cache_new_con {
-                        LAST_QUERY.store(
-                            chrono::Utc::now().timestamp().try_into().unwrap_or(0),
-                            std::sync::atomic::Ordering::Relaxed,
-                        );
-                        tokio::spawn(async move {
-                            loop {
-                                tokio::time::sleep(Duration::from_secs(5)).await;
-                                let last_query =
-                                    LAST_QUERY.load(std::sync::atomic::Ordering::Relaxed);
-                                let now = chrono::Utc::now().timestamp().try_into().unwrap_or(0);
-
-                                //we cache connection for 5 minutes at most
-                                if last_query + 60 * 1 < now {
-                                    // tracing::error!("Closing cache connection due to inactivity");
-                                    tracing::info!(
-                                        "Closing cache pg executor connection due to inactivity"
-                                    );
-                                    break;
-                                }
-                                let mtex = CONNECTION_CACHE.lock().await;
-                                if mtex.is_none() {
-                                    // connection is not in the mutex anymore
-                                    break;
-                                } else if let Some(mtex) = mtex.as_ref() {
-                                    if mtex.0.as_str() != &database_string_clone {
-                                        // connection is not the latest one
-                                        break;
-                                    }
-                                }
-
-                                tracing::debug!(
-                                    "Keeping cached pg executor connection alive due to activity"
-                                )
-                            }
-                            let mut mtex = CONNECTION_CACHE.lock().await;
-                            *mtex = None;
-                            abort_handler.abort();
-                        });
-                    }
-                } else {
-                    handle.abort();
-                }
-            } else {
-                handle.abort();
-            }
-        } else {
-            handle.abort();
-        }
+    if !*CLOUD_HOSTED && lease.slot == PgLeaseSlot::Uncached {
+        lease.cache_on_release = is_most_used_conn(&database_string).await;
     }
+    drop(lease);
+
     *mem_peak = (result.get().len() / 1000) as i32;
     // And then check that we got back the same string we sent over.
     return Ok(result);