diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index df07aba681..084380747e 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -29,6 +29,7 @@ use tokio_postgres::{AsyncMessage, NoTls, Row, SimpleQueryMessage, Socket}; use tokio_util::sync::CancellationToken; use super::file_validator::validate_file_path; +use crate::models::connection::DatabaseType; use crate::query::{await_stream_with_progress_timeout, DbOperationBudget, StreamProgressClock}; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ @@ -7765,6 +7766,7 @@ pub async fn execute_query_in_read_only_transaction_with_rollback( pub async fn stream_select_query_with_cancel( pool: &Pool, + db_type: Option, schema: Option<&str>, setup_sql: &[String], sql: &str, @@ -7861,7 +7863,7 @@ pub async fn stream_select_query_with_cancel( }; if schema_was_set { - let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; + let reset_result = reset_postgres_search_path(&client, db_type, budget.cleanup_timeout, start).await; match (result, reset_result) { (Ok(rows), Ok(())) => Ok(rows), (Err(query_err), Ok(())) => Err(query_err), @@ -7874,11 +7876,12 @@ pub async fn stream_select_query_with_cancel( } pub async fn execute_query_with_schema(pool: &Pool, schema: &str, sql: &str) -> Result { - execute_query_with_schema_and_max_rows(pool, schema, sql, None).await + execute_query_with_schema_and_max_rows(pool, None, schema, sql, None).await } pub async fn execute_query_with_schema_and_max_rows( pool: &Pool, + db_type: Option, schema: &str, sql: &str, max_rows: Option, @@ -7920,12 +7923,13 @@ pub async fn execute_query_with_schema_and_max_rows( result.is_ok() ); - let reset_result = reset_postgres_search_path(&client, super::connection_timeout(), start).await; + let reset_result = reset_postgres_search_path(&client, db_type, super::connection_timeout(), start).await; merge_postgres_query_and_reset_result(result, reset_result) } pub async fn execute_query_with_schema_and_max_rows_and_cancel( pool: &Pool, + db_type: Option, schema: &str, sql: &str, max_rows: Option, @@ -7991,17 +7995,31 @@ pub async fn execute_query_with_schema_and_max_rows_and_cancel( result.is_ok() ); - let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; + let reset_result = reset_postgres_search_path(&client, db_type, budget.cleanup_timeout, start).await; merge_postgres_query_and_reset_result(result, reset_result) } +/// GaussDB/openGauss reject PostgreSQL's `RESET search_path` syntax, so the +/// post-query schema cleanup must re-issue `SET search_path TO DEFAULT` for +/// those engines. Every other backend — and an unknown (`None`) type — keeps +/// the historical `RESET search_path` behavior. +pub(crate) fn reset_search_path_sql(db_type: Option) -> &'static str { + match db_type { + Some(DatabaseType::Gaussdb | DatabaseType::OpenGauss) => "SET search_path TO DEFAULT", + _ => "RESET search_path", + } +} + async fn reset_postgres_search_path( client: &deadpool_postgres::Client, + db_type: Option, timeout_duration: Duration, start: Instant, ) -> Result<(), String> { let reset_start = Instant::now(); - match execute_postgres_infra_statement(client, "RESET search_path", timeout_duration, "schema.reset").await { + match execute_postgres_infra_statement(client, reset_search_path_sql(db_type), timeout_duration, "schema.reset") + .await + { Ok(_) => { log::info!( "[postgres][execute_with_schema:reset-search-path:done] elapsed_ms={} total_ms={}", @@ -10559,6 +10577,14 @@ mod tests { ); } + #[test] + fn postgres_reset_search_path_sql_selects_dialect_compatible_statement() { + assert_eq!(reset_search_path_sql(Some(DatabaseType::Gaussdb)), "SET search_path TO DEFAULT"); + assert_eq!(reset_search_path_sql(Some(DatabaseType::OpenGauss)), "SET search_path TO DEFAULT"); + assert_eq!(reset_search_path_sql(Some(DatabaseType::Postgres)), "RESET search_path"); + assert_eq!(reset_search_path_sql(None), "RESET search_path"); + } + #[test] fn postgres_single_schema_search_path_preserves_scope_and_quoting() { assert_eq!( @@ -13893,6 +13919,7 @@ mod tests { let mut streamed_rows = Vec::new(); let streaming_result = stream_select_query_with_cancel( &pool, + None, Some(&schema), &[], &query_sql, diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 60c3eecc96..8a552684b1 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -2073,6 +2073,7 @@ async fn do_execute_typed( let cancel_context = state.get_postgres_cancel_context(pool_key).await; let result = execute_postgres_pool_statement( &p, + pool_db_type, schema.as_deref(), sql, max_rows, @@ -2092,6 +2093,7 @@ async fn do_execute_typed( ); execute_postgres_pool_statement( &p, + pool_db_type, schema.as_deref(), &fallback_sql, max_rows, @@ -2574,6 +2576,7 @@ fn postgres_preview_fallback_retry_sql(options: &QueryExecutionOptions, error: & #[allow(clippy::too_many_arguments)] async fn execute_postgres_pool_statement( pool: &deadpool_postgres::Pool, + db_type: Option, schema: Option<&str>, sql: &str, max_rows: Option, @@ -2597,6 +2600,7 @@ async fn execute_postgres_pool_statement( } else if let Some(schema) = schema { db::postgres::execute_query_with_schema_and_max_rows_and_cancel( pool, + db_type, schema, sql, max_rows, @@ -4747,7 +4751,7 @@ pub async fn execute_statements_in_transaction_on_pool_typed( let result = match path { Some(BatchTransactionPath::Pg(pool)) => { let cancel_context = state.get_postgres_cancel_context(pool_key).await; - exec_tx_pg_inner(pool, statements, schema, start, operation_budget.clone(), cancel_context).await + exec_tx_pg_inner(pool, db_type, statements, schema, start, operation_budget.clone(), cancel_context).await } Some(BatchTransactionPath::Mysql(pool)) => exec_tx_mysql_inner( state, @@ -4870,6 +4874,7 @@ fn batch_transaction_path(pool: &PoolKind) -> BatchTransactionPath { async fn exec_tx_pg_inner( pool: deadpool_postgres::Pool, + db_type: Option, statements: &[String], schema: Option<&str>, start: std::time::Instant, @@ -4892,11 +4897,11 @@ async fn exec_tx_pg_inner( } let tx_result = exec_tx_pg_statements(&mut client, statements, &budget, cancel_context).await; - // Always reset search_path so the connection is clean when returned to the pool + // GaussDB/openGauss reject PostgreSQL's RESET search_path syntax. let reset_result = if had_schema { db::postgres::execute_postgres_infra_statement( &client, - "RESET search_path", + db::postgres::reset_search_path_sql(db_type), budget.cleanup_timeout, "schema.reset", ) diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 6a268e4275..8b0b0a07d0 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -1173,6 +1173,7 @@ async fn try_export_postgres_query_result_stream( let setup_sql = safe_postgres_temp_setup_sql(&request.setup_sql).unwrap_or_default(); let stream_result = crate::db::postgres::stream_select_query_with_cancel( &pool, + Some(request.database_type), request.schema.as_deref(), &setup_sql, &request.sql,