fix: datatable full schema hangs behind a transaction-pooling postgres proxy (#10352)

* fix: datatable full schema hangs behind a transaction-pooling postgres proxy

* test: pause the clock in the pg connection shutdown test

backend/windmill-api-settings/src/lib.rs

@@ -1727,12 +1727,7 @@ async fn setup_custom_instance_pg_database_inner(
     logs.grant_permissions = "OK".to_string();
 
     drop(client); // /!\ Drop before joining to avoid deadlock
-    join_handle
-        .await
-        .map_err(|e| error::Error::ExecutionErr(format!("join error: {}", e.to_string())))?
-        .map_err(|e| {
-            error::Error::ExecutionErr(format!("tokio_postgres error: {}", e.to_string()))
-        })?;
+    windmill_common::shutdown_pg_connection(join_handle).await?;
 
     Ok(())
 }

backend/windmill-api-workspaces/src/workspaces.rs

@@ -2677,7 +2677,7 @@ async fn create_pg_database(
 
         if db_exists {
             drop(client);
-            let _ = join_handle.await;
+            let _ = windmill_common::shutdown_pg_connection(join_handle).await;
             return Err(Error::BadRequest(format!(
                 "Database '{}' already exists on the resource server",
                 req.target_dbname
@@ -2695,10 +2695,7 @@ async fn create_pg_database(
             })?;
 
         drop(client);
-        join_handle
-            .await
-            .map_err(|e| Error::internal_err(format!("join error: {}", e)))?
-            .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?;
+        windmill_common::shutdown_pg_connection(join_handle).await?;
     }
 
     Ok(format!("Created database '{}'", req.target_dbname))
@@ -2801,10 +2798,7 @@ async fn get_datatable_full_schema(
         .map_err(Error::internal_err)?;
 
     drop(client);
-    join_handle
-        .await
-        .map_err(|e| Error::internal_err(format!("join error: {}", e)))?
-        .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?;
+    windmill_common::shutdown_pg_connection(join_handle).await?;
 
     Ok(Json(result))
 }
@@ -6389,10 +6383,7 @@ async fn snapshot_datatable_schema(
         .map_err(Error::internal_err)?;
 
     drop(client);
-    join_handle
-        .await
-        .map_err(|e| Error::internal_err(format!("join error: {}", e)))?
-        .map_err(|e| Error::internal_err(format!("tokio_postgres error: {}", e)))?;
+    windmill_common::shutdown_pg_connection(join_handle).await?;
 
     serde_json::to_value(schema)
         .map_err(|e| Error::internal_err(format!("Failed to serialize schema: {}", e)))

backend/windmill-api-workspaces/src/workspaces_extra.rs

@@ -1352,7 +1352,7 @@ pub async fn drop_forked_datatable_databases(
                         ));
                     }
                     drop(client);
-                    let _ = join_handle.await;
+                    let _ = windmill_common::shutdown_pg_connection(join_handle).await;
                 }
                 Err(e) => {
                     errors.push(format!(
@@ -1681,7 +1681,7 @@ async fn drop_fork_ducklake_metadata_schema(
         )
         .await;
     drop(client);
-    let _ = join_handle.await;
+    let _ = windmill_common::shutdown_pg_connection(join_handle).await;
     res.map_err(|e| Error::internal_err(format!("{e:#}")))?;
     Ok(())
 }

backend/windmill-common/Cargo.toml

@@ -130,6 +130,10 @@ opentelemetry-appender-tracing = { workspace = true, optional = true }
 tonic = { workspace = true, optional = true }
 equivalent = "1.0.2"
 
+[dev-dependencies]
+# `test-util` is not part of tokio's `full`; it is what lets tests pause the clock.
+tokio = { workspace = true, features = ["test-util"] }
+
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
 tikv-jemalloc-ctl = { optional = true, workspace = true }
 

backend/windmill-common/src/lib.rs

@@ -35,6 +35,7 @@ pub mod auth;
 pub mod bench;
 pub mod cache;
 pub mod client;
+pub mod data_metrics;
 pub mod db;
 #[cfg(all(feature = "enterprise", feature = "private"))]
 mod db_entra_ee;
@@ -62,7 +63,6 @@ pub mod instance_config;
 pub mod job_metrics;
 pub mod log_context;
 pub mod materialization;
-pub mod data_metrics;
 pub mod min_version;
 pub mod notify_events;
 pub mod runtime_assets;
@@ -1068,6 +1068,49 @@ impl PgDatabase {
     }
 }
 
+/// How long a `tokio_postgres` connection task gets to wind down once its `Client` is dropped.
+const PG_CONNECTION_SHUTDOWN_GRACE: std::time::Duration = std::time::Duration::from_secs(5);
+
+/// Wind down the task driving a `tokio_postgres` connection after its `Client` has been dropped,
+/// surfacing whatever error the connection ended with. A teardown that has to be aborted is
+/// reported as success — the work the client did is already done and complete.
+///
+/// The task only finishes once the exchange the client left behind (its Terminate, and any
+/// still-unanswered request) has been settled by the peer. A connection proxy that stops
+/// replying leaves that pending forever, so waiting on the task without a deadline pins the
+/// caller and the socket for the lifetime of the process. Aborting past the grace period drops
+/// the stream, which is the only cleanup the task owes.
+pub async fn shutdown_pg_connection(
+    join_handle: tokio::task::JoinHandle<Result<(), tokio_postgres::Error>>,
+) -> error::Result<()> {
+    let abort_handle = join_handle.abort_handle();
+    match tokio::time::timeout(PG_CONNECTION_SHUTDOWN_GRACE, join_handle).await {
+        Ok(Ok(Ok(()))) => Ok(()),
+        Ok(Ok(Err(e))) => Err(error::Error::internal_err(format!(
+            "tokio_postgres error: {}",
+            e
+        ))),
+        Ok(Err(e)) => Err(error::Error::internal_err(format!("join error: {}", e))),
+        Err(_) => {
+            tracing::warn!(
+                "Postgres connection did not close within {}s of its client being dropped, aborting it",
+                PG_CONNECTION_SHUTDOWN_GRACE.as_secs()
+            );
+            abort_handle.abort();
+            Ok(())
+        }
+    }
+}
+
+#[cfg(test)]
+mod pg_connection_shutdown_tests {
+    #[tokio::test(start_paused = true)]
+    async fn gives_up_on_a_connection_task_that_never_finishes() {
+        let never_finishes = tokio::spawn(std::future::pending());
+        assert!(super::shutdown_pg_connection(never_finishes).await.is_ok());
+    }
+}
+
 /// Validate a database name to prevent SQL injection.
 /// Must start with a letter, contain only alphanumeric characters, underscores, or hyphens, and be <= 63 chars.
 pub fn validate_dbname(dbname: &str) -> error::Result<()> {
@@ -1219,10 +1262,7 @@ pub async fn create_custom_instance_database(
     }
 
     drop(client);
-    join_handle
-        .await
-        .map_err(|e| error::Error::internal_err(format!("join error: {}", e)))?
-        .map_err(|e| error::Error::internal_err(format!("tokio_postgres error: {}", e)))?;
+    shutdown_pg_connection(join_handle).await?;
 
     // Register in global_settings
     let status_json = serde_json::json!({

backend/windmill-common/src/query_builders.rs

@@ -4830,8 +4830,37 @@ fn pg_action_to_string(action: &str) -> String {
     }
 }
 
+/// Rows of a simple-protocol result, dropping the framing messages.
+fn simple_query_rows(
+    messages: Vec<tokio_postgres::SimpleQueryMessage>,
+) -> Vec<tokio_postgres::SimpleQueryRow> {
+    messages
+        .into_iter()
+        .filter_map(|m| match m {
+            tokio_postgres::SimpleQueryMessage::Row(row) => Some(row),
+            _ => None,
+        })
+        .collect()
+}
+
+fn required_str<'a>(
+    row: &'a tokio_postgres::SimpleQueryRow,
+    column: &str,
+) -> Result<&'a str, String> {
+    row.try_get(column)
+        .map_err(|e| format!("Failed to read column {}: {}", column, e))?
+        .ok_or_else(|| format!("Unexpected NULL in column {}", column))
+}
+
 /// Introspect a PostgreSQL database and return the full schema.
 /// Takes a connected tokio_postgres Client.
+///
+/// Both statements go through the simple query protocol. The extended protocol allocates a
+/// named prepared statement per call and closes it when the statement handle drops; behind a
+/// transaction-pooling proxy those names are shared with, and outlive, other sessions on the
+/// same backend, and the exchange then stalls with no reply — the connection never becomes
+/// idle again and the request hangs. Neither statement takes parameters, so nothing here
+/// needs the extended protocol.
 pub async fn pg_get_full_schema(
     client: &tokio_postgres::Client,
 ) -> Result<FullDatabaseSchema, String> {
@@ -4840,7 +4869,7 @@ pub async fn pg_get_full_schema(
     // per-column correlated subqueries — on large catalogs those subqueries run
     // once per column and make the introspection time out.
     let column_rows = client
-        .query(
+        .simple_query(
             "SELECT
                 ns.nspname AS schema_name,
                 c.relname AS table_name,
@@ -4862,13 +4891,13 @@ pub async fn pg_get_full_schema(
                 AND NOT a.attisdropped
                 AND ns.nspname NOT IN ('pg_catalog', 'information_schema')
             ORDER BY ns.nspname, c.relname, a.attnum",
-            &[],
         )
         .await
+        .map(simple_query_rows)
         .map_err(|e| format!("Failed to query columns: {}", e))?;
 
     let fk_rows = client
-        .query(
+        .simple_query(
             "SELECT
                 ns.nspname AS schema_name,
                 c.relname AS table_name,
@@ -4890,21 +4919,21 @@ pub async fn pg_get_full_schema(
             WHERE con.contype = 'f'
                 AND ns.nspname NOT IN ('pg_catalog', 'information_schema')
             ORDER BY ns.nspname, c.relname, con.conname, u.ord",
-            &[],
         )
         .await
+        .map(simple_query_rows)
         .map_err(|e| format!("Failed to query foreign keys: {}", e))?;
 
     let mut result: FullDatabaseSchema = std::collections::HashMap::new();
 
     for row in &column_rows {
-        let schema_name: &str = row.get("schema_name");
-        let table_name: &str = row.get("table_name");
-        let column_name: &str = row.get("column_name");
-        let datatype: &str = row.get("datatype");
+        let schema_name = required_str(row, "schema_name")?;
+        let table_name = required_str(row, "table_name")?;
+        let column_name = required_str(row, "column_name")?;
+        let datatype = required_str(row, "datatype")?;
         let default_value: Option<&str> = row.get("default_value");
-        let nullable: bool = row.get("nullable");
-        let is_primary_key: bool = row.get("is_primary_key");
+        let nullable = required_str(row, "nullable")? == "t";
+        let is_primary_key = required_str(row, "is_primary_key")? == "t";
         let pk_constraint_name: Option<&str> = row.get("pk_constraint_name");
 
         let schema_tables = result.entry(schema_name.to_string()).or_default();
@@ -4937,15 +4966,15 @@ pub async fn pg_get_full_schema(
     > = std::collections::HashMap::new();
 
     for row in &fk_rows {
-        let schema_name: &str = row.get("schema_name");
-        let table_name: &str = row.get("table_name");
-        let fk_name: &str = row.get("fk_constraint_name");
-        let source_column: &str = row.get("source_column");
-        let ref_schema: &str = row.get("ref_schema");
-        let ref_table: &str = row.get("ref_table");
-        let ref_column: &str = row.get("ref_column");
-        let on_delete: &str = row.get("on_delete");
-        let on_update: &str = row.get("on_update");
+        let schema_name = required_str(row, "schema_name")?;
+        let table_name = required_str(row, "table_name")?;
+        let fk_name = required_str(row, "fk_constraint_name")?;
+        let source_column = required_str(row, "source_column")?;
+        let ref_schema = required_str(row, "ref_schema")?;
+        let ref_table = required_str(row, "ref_table")?;
+        let ref_column = required_str(row, "ref_column")?;
+        let on_delete = required_str(row, "on_delete")?;
+        let on_update = required_str(row, "on_update")?;
 
         let target_table = if ref_schema == schema_name {
             ref_table.to_string()

backend/windmill-common/src/workspaces.rs

@@ -1866,7 +1866,7 @@ async fn inspect_fork_catalog(
     );
     let table_rows = client.query(&qt, &[]).await;
     drop(client);
-    let _ = join_handle.await;
+    let _ = crate::shutdown_pg_connection(join_handle).await;
 
     existing_schemas.extend(
         same_catalog_res
@@ -1920,7 +1920,7 @@ async fn inspect_fork_catalog(
             let join_handle = tokio::spawn(async move { connection.await });
             let res = query_schemas(&client, schemas).await;
             drop(client);
-            let _ = join_handle.await;
+            let _ = crate::shutdown_pg_connection(join_handle).await;
             res.map_err(|e| Error::internal_err(format!("{e}")))
         }
         .await;