diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 333c1ee9c..2659ad1ad 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -463,6 +463,19 @@ pub async fn execute_query_with_schema_and_max_rows( .await .map_err(|e| e.to_string())?; + let result = execute_query_with_max_rows_inner(&client, sql, max_rows).await; + + // Always reset search_path so the connection is clean when returned to the pool + let _ = client.execute("RESET search_path", &[]).await; + + result +} + +async fn execute_query_with_max_rows_inner( + client: &deadpool_postgres::Client, + sql: &str, + max_rows: Option, +) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index f2f68b3cc..9fec9f2a9 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -933,12 +933,35 @@ async fn exec_tx_pg_inner( start: std::time::Instant, ) -> Result { let mut client = pool.get().await.map_err(|e| format!("Failed to acquire connection: {}", e))?; + let had_schema = schema.is_some(); if let Some(s) = schema { client .execute(&format!("SET search_path TO {}, public", db::postgres::pg_quote_ident(s)), &[]) .await .map_err(|e| format!("SET search_path failed: {}", e))?; } + let tx_result = exec_tx_pg_statements(&mut client, statements).await; + + // Always reset search_path so the connection is clean when returned to the pool + if had_schema { + let _ = client.execute("RESET search_path", &[]).await; + } + + match tx_result { + Ok(total_affected) => Ok(db::QueryResult { + columns: vec![], + rows: vec![], + affected_rows: total_affected, + execution_time_ms: start.elapsed().as_millis(), + truncated: false, + session_id: None, + has_more: false, + }), + Err(e) => Err(e), + } +} + +async fn exec_tx_pg_statements(client: &mut deadpool_postgres::Client, statements: &[String]) -> Result { let tx = client.transaction().await.map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { @@ -951,15 +974,7 @@ async fn exec_tx_pg_inner( } } tx.commit().await.map_err(|e| format!("COMMIT failed: {}", e))?; - Ok(db::QueryResult { - columns: vec![], - rows: vec![], - affected_rows: total_affected, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - session_id: None, - has_more: false, - }) + Ok(total_affected) } async fn exec_tx_mysql_inner(