From 68ad9b2c5d045e25c5c107570bd2f8fddc28cd7a Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 19 Aug 2026 11:44:40 +0800 Subject: [PATCH] fix(postgres): recover invalidated prepared statements --- Cargo.lock | 6 +- Cargo.toml | 6 +- agents/drivers/duckdb/Cargo.lock | 6 +- agents/drivers/duckdb/Cargo.toml | 6 +- crates/dbx-core/src/db/postgres.rs | 485 +++++++++++++++--- crates/dbx-core/src/query.rs | 12 +- .../dbx-core/tests/live_manual_transaction.rs | 40 +- 7 files changed, 478 insertions(+), 83 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fa06cf0ac3..da4c4b239f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6151,7 +6151,7 @@ dependencies = [ [[package]] name = "postgres-protocol" version = "0.6.12" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "base64 0.23.0", "byteorder", @@ -6170,7 +6170,7 @@ dependencies = [ [[package]] name = "postgres-types" version = "0.2.14" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "bytes", "chrono", @@ -9101,7 +9101,7 @@ dependencies = [ [[package]] name = "tokio-postgres" version = "0.7.18" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "async-trait", "byteorder", diff --git a/Cargo.toml b/Cargo.toml index 921a639106..81a3d88db6 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,9 +19,9 @@ pageant = { path = "vendor/pageant" } # F6 "Focus Next Pane" accelerator to prevent a black screen on frameless # Overlay-titlebar windows (see vendor/wry/src/webview2/mod.rs). wry = { path = "vendor/wry" } -tokio-postgres = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } -postgres-types = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } -postgres-protocol = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } +tokio-postgres = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } +postgres-types = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } +postgres-protocol = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } # Temporary compatibility patch for MySQL accounts using the deprecated # sha256_password auth plugin. Remove after upstream support lands. mysql_async = { git = "https://github.com/t8y2/mysql_async.git", rev = "b615bb0b10b5975182a4aa069cabc3d09021f6ab" } diff --git a/agents/drivers/duckdb/Cargo.lock b/agents/drivers/duckdb/Cargo.lock index b534a02737..4c99eca22b 100644 --- a/agents/drivers/duckdb/Cargo.lock +++ b/agents/drivers/duckdb/Cargo.lock @@ -3682,7 +3682,7 @@ dependencies = [ [[package]] name = "postgres-protocol" version = "0.6.12" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "base64 0.23.0", "byteorder", @@ -3701,7 +3701,7 @@ dependencies = [ [[package]] name = "postgres-types" version = "0.2.14" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "bytes", "chrono", @@ -5479,7 +5479,7 @@ dependencies = [ [[package]] name = "tokio-postgres" version = "0.7.18" -source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=115f9fef10f0fc3669b5337955e4eb461fc349a6#115f9fef10f0fc3669b5337955e4eb461fc349a6" +source = "git+https://github.com/t8y2/tokio-postgres-gaussdb.git?rev=ac8ad476de59093d33ab9f294531c956e0c6a724#ac8ad476de59093d33ab9f294531c956e0c6a724" dependencies = [ "async-trait", "byteorder", diff --git a/agents/drivers/duckdb/Cargo.toml b/agents/drivers/duckdb/Cargo.toml index 61f5c72a0e..eaf36a19e3 100644 --- a/agents/drivers/duckdb/Cargo.toml +++ b/agents/drivers/duckdb/Cargo.toml @@ -37,9 +37,9 @@ tokio-util = "0.7" uuid = { version = "1", features = ["v4"] } [patch.crates-io] -tokio-postgres = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } -postgres-types = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } -postgres-protocol = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "115f9fef10f0fc3669b5337955e4eb461fc349a6" } +tokio-postgres = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } +postgres-types = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } +postgres-protocol = { git = "https://github.com/t8y2/tokio-postgres-gaussdb.git", rev = "ac8ad476de59093d33ab9f294531c956e0c6a724" } mysql_async = { git = "https://github.com/t8y2/mysql_async.git", rev = "b615bb0b10b5975182a4aa069cabc3d09021f6ab" } [profile.release] diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 36a05724c6..e01366b035 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1083,14 +1083,92 @@ fn should_retry_postgres_stale_cache_fields(sqlstate: Option<&str>, routine: Opt structured_match || message.to_ascii_lowercase().contains("cached plan must not change result type") } +fn should_fallback_postgres_missing_prepared_statement(err: &tokio_postgres::Error) -> bool { + if let Some(db_error) = err.as_db_error() { + return should_fallback_postgres_missing_prepared_statement_fields( + Some(db_error.code().code()), + db_error.message(), + ); + } + should_fallback_postgres_missing_prepared_statement_fields(None, &err.to_string()) +} + +fn should_fallback_postgres_missing_prepared_statement_fields(sqlstate: Option<&str>, message: &str) -> bool { + if sqlstate == Some("26000") { + return true; + } + let message = message.to_ascii_lowercase(); + message.contains("prepared statement") && message.contains("does not exist") +} + +fn postgres_unnamed_statement_clients() -> &'static Mutex>> { + static CLIENTS: OnceLock>>> = OnceLock::new(); + CLIENTS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn postgres_client_uses_unnamed_statements(client: &deadpool_postgres::Client) -> bool { + let statement_cache = &client.statement_cache; + let key = Arc::as_ptr(statement_cache) as usize; + let mut clients = postgres_unnamed_statement_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + match clients.get(&key).and_then(Weak::upgrade) { + Some(cached) if Arc::ptr_eq(&cached, statement_cache) => true, + _ => { + clients.remove(&key); + false + } + } +} + +fn mark_postgres_client_unnamed_statements(client: &deadpool_postgres::Client) { + let statement_cache = &client.statement_cache; + let key = Arc::as_ptr(statement_cache) as usize; + let mut clients = postgres_unnamed_statement_clients().lock().unwrap_or_else(|poisoned| poisoned.into_inner()); + clients.retain(|_, cached| cached.strong_count() > 0); + clients.insert(key, Arc::downgrade(statement_cache)); +} + +fn postgres_typed_params<'a>( + params: &[&'a (dyn tokio_postgres::types::ToSql + Sync)], + param_types: &[Type], +) -> Option> { + (params.len() == param_types.len()).then(|| params.iter().copied().zip(param_types.iter().cloned()).collect()) +} + +async fn postgres_query_unnamed( + client: &deadpool_postgres::Client, + sql: &str, +) -> Result { + client.query_typed_raw(sql, std::iter::empty::<(&(dyn tokio_postgres::types::ToSql + Sync), Type)>()).await +} + async fn postgres_query_cached( client: &deadpool_postgres::Client, sql: &str, params: &[&(dyn tokio_postgres::types::ToSql + Sync)], ) -> Result, tokio_postgres::Error> { + if postgres_client_uses_unnamed_statements(client) && params.is_empty() { + return client.query_typed(sql, &[]).await; + } let stmt = client.prepare_cached(sql).await?; + if postgres_client_uses_unnamed_statements(client) { + let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { + return client.query(&stmt, params).await; + }; + return client.query_typed(sql, &typed_params).await; + } match client.query(&stmt, params).await { Ok(rows) => Ok(rows), + Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { + let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { + return Err(err); + }; + log::warn!( + "[postgres][metadata:missing_prepared_statement] downgrading connection to unnamed statements: {}", + pg_error_to_string(err) + ); + mark_postgres_client_unnamed_statements(client); + client.query_typed(sql, &typed_params).await + } Err(err) if should_retry_postgres_stale_cache(&err) => { // Metadata queries can be cached while a table/view definition is // changed from another session. Evict and retry once with fresh @@ -1111,9 +1189,29 @@ async fn postgres_query_one_cached( sql: &str, params: &[&(dyn tokio_postgres::types::ToSql + Sync)], ) -> Result { + if postgres_client_uses_unnamed_statements(client) && params.is_empty() { + return client.query_typed_one(sql, &[]).await; + } let stmt = client.prepare_cached(sql).await?; + if postgres_client_uses_unnamed_statements(client) { + let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { + return client.query_one(&stmt, params).await; + }; + return client.query_typed_one(sql, &typed_params).await; + } match client.query_one(&stmt, params).await { Ok(row) => Ok(row), + Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { + let Some(typed_params) = postgres_typed_params(params, stmt.params()) else { + return Err(err); + }; + log::warn!( + "[postgres][metadata_one:missing_prepared_statement] downgrading connection to unnamed statements: {}", + pg_error_to_string(err) + ); + mark_postgres_client_unnamed_statements(client); + client.query_typed_one(sql, &typed_params).await + } Err(err) if should_retry_postgres_stale_cache(&err) => { // Same stale-cache protection as postgres_query_cached, for scalar // catalog probes such as pg_proc feature detection. @@ -1140,15 +1238,15 @@ struct PreparedSelectMetadata { unsupported_type: Option, } -fn prepared_select_metadata(stmt: &tokio_postgres::Statement) -> PreparedSelectMetadata { - let columns: Vec = stmt.columns().iter().map(|c| c.name().to_string()).collect(); - let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); +fn prepared_select_metadata(columns: &[tokio_postgres::Column]) -> PreparedSelectMetadata { + let column_names: Vec = columns.iter().map(|c| c.name().to_string()).collect(); + let column_types: Vec = columns.iter().map(|c| c.type_().name().to_string()).collect(); let column_classes = classify_pg_column_types(&column_types); - let unsupported_type = stmt.columns().iter().zip(&column_classes).find_map(|(column, col_type)| { + let unsupported_type = columns.iter().zip(&column_classes).find_map(|(column, col_type)| { let pg_type = column.type_(); pg_type_requires_text_protocol(pg_type, *col_type).then(|| pg_type.name().to_string()) }); - PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type } + PreparedSelectMetadata { columns: column_names, column_types, column_classes, unsupported_type } } async fn prepare_select_with_metadata( @@ -1156,45 +1254,88 @@ async fn prepare_select_with_metadata( sql: &str, ) -> Result<(tokio_postgres::Statement, PreparedSelectMetadata), tokio_postgres::Error> { let mut stmt = client.prepare_cached(sql).await?; - let mut metadata = prepared_select_metadata(&stmt); + let mut metadata = prepared_select_metadata(stmt.columns()); if metadata.unsupported_type.is_some() { stmt = client.prepare(sql).await?; - metadata = prepared_select_metadata(&stmt); + metadata = prepared_select_metadata(stmt.columns()); } Ok((stmt, metadata)) } +enum PostgresSelectStreamOutcome { + Binary { stream: tokio_postgres::RowStream, metadata: PreparedSelectMetadata }, + TextFallback { column_types: Vec, unsupported_type: String }, +} + +fn postgres_select_stream_outcome(stream: tokio_postgres::RowStream) -> PostgresSelectStreamOutcome { + let metadata = prepared_select_metadata(stream.columns()); + if let Some(unsupported_type) = metadata.unsupported_type.clone() { + return PostgresSelectStreamOutcome::TextFallback { column_types: metadata.column_types, unsupported_type }; + } + PostgresSelectStreamOutcome::Binary { stream, metadata } +} + +async fn start_postgres_select_stream( + client: &deadpool_postgres::Client, + sql: &str, + force_unnamed: bool, +) -> Result { + if force_unnamed || postgres_client_uses_unnamed_statements(client) { + return postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome); + } + + let (stmt, metadata) = prepare_select_with_metadata(client, sql).await?; + if let Some(unsupported_type) = metadata.unsupported_type { + return Ok(PostgresSelectStreamOutcome::TextFallback { column_types: metadata.column_types, unsupported_type }); + } + + let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); + match client.query_raw(&stmt, params).await { + Ok(stream) => Ok(postgres_select_stream_outcome(stream)), + Err(err) if should_fallback_postgres_missing_prepared_statement(&err) => { + // Bind failed before execution. Downgrade this physical connection + // so later queries do not repeat the named-statement round trip. + log::warn!( + "[postgres][prepared_statement:missing] downgrading connection to unnamed statements: {}", + pg_error_to_string(err) + ); + mark_postgres_client_unnamed_statements(client); + postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome) + } + Err(err) => Err(err), + } +} + async fn execute_select_prepared( client: &deadpool_postgres::Client, sql: &str, start: Instant, row_limit: usize, progress_clock: Option<&StreamProgressClock>, + force_unnamed: bool, ) -> Result { - let prepared_start = Instant::now(); - let (stmt, metadata) = prepare_select_with_metadata(client, sql).await?; + let stream_start = Instant::now(); + let stream_outcome = start_postgres_select_stream(client, sql, force_unnamed).await?; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } log::info!( - "[postgres][select:prepare_cached:done] elapsed_ms={} total_ms={}", - prepared_start.elapsed().as_millis(), + "[postgres][select:stream_ready] elapsed_ms={} total_ms={}", + stream_start.elapsed().as_millis(), start.elapsed().as_millis() ); - let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type } = metadata; - if let Some(unsupported_type) = unsupported_type { - return Ok(PreparedSelectOutcome::TextFallback { column_types, unsupported_type }); - } - - let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); - let query_start = Instant::now(); - let stream = client.query_raw(&stmt, params).await?; + let (stream, metadata) = match stream_outcome { + PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), + PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { + return Ok(PreparedSelectOutcome::TextFallback { column_types, unsupported_type }); + } + }; + let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type: _ } = metadata; if let Some(progress_clock) = progress_clock { progress_clock.mark(); } log::info!( - "[postgres][select:query_raw:done] elapsed_ms={} total_ms={} column_count={}", - query_start.elapsed().as_millis(), + "[postgres][select:metadata_ready] total_ms={} column_count={}", start.elapsed().as_millis(), columns.len() ); @@ -1380,7 +1521,16 @@ pub(crate) async fn execute_select_query( start: Instant, row_limit: usize, ) -> Result { - execute_select_query_with_progress(client, sql, start, row_limit, None).await + execute_select_query_with_progress(client, sql, start, row_limit, None, false).await +} + +pub(crate) async fn execute_select_query_unnamed( + client: &deadpool_postgres::Client, + sql: &str, + start: Instant, + row_limit: usize, +) -> Result { + execute_select_query_with_progress(client, sql, start, row_limit, None, true).await } async fn execute_select_query_with_progress( @@ -1389,8 +1539,9 @@ async fn execute_select_query_with_progress( start: Instant, row_limit: usize, progress_clock: Option<&StreamProgressClock>, + force_unnamed: bool, ) -> Result { - match execute_select_prepared(client, sql, start, row_limit, progress_clock).await { + match execute_select_prepared(client, sql, start, row_limit, progress_clock, force_unnamed).await { Ok(outcome) => finish_prepared_select(client, sql, start, row_limit, outcome, progress_clock).await, Err(err) if should_retry_postgres_stale_cache(&err) => { // The cached prepared statement is stale (e.g. the view or table @@ -1398,7 +1549,7 @@ async fn execute_select_query_with_progress( // stale entry and retry with a fresh server-side prepare. log::warn!("[postgres][select:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); - match execute_select_prepared(client, sql, start, row_limit, progress_clock).await { + match execute_select_prepared(client, sql, start, row_limit, progress_clock, force_unnamed).await { Ok(outcome) => finish_prepared_select(client, sql, start, row_limit, outcome, progress_clock).await, Err(err) if should_retry_postgres_text_query(&err) => { execute_select_text(client, sql, start, row_limit, None, progress_clock).await @@ -1441,20 +1592,18 @@ async fn stream_select_query_prepared( sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, + force_unnamed: bool, ) -> Result { - let (stmt, metadata) = prepare_select_with_metadata(client, sql) - .await - .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; - let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type } = metadata; - if let Some(unsupported_type) = unsupported_type { - return Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }); - } - - let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); - let stream = client - .query_raw(&stmt, params) + let stream_outcome = start_postgres_select_stream(client, sql, force_unnamed) .await .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; + let (stream, metadata) = match stream_outcome { + PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), + PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { + return Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }); + } + }; + let PreparedSelectMetadata { columns, column_types, column_classes, unsupported_type: _ } = metadata; tokio::pin!(stream); let mut rows_streamed = 0_u64; let mut columns_emitted = false; @@ -1528,13 +1677,23 @@ async fn stream_select_query_text( Ok(rows_streamed) } -pub(crate) async fn stream_select_query_inner( +pub(crate) async fn stream_select_query_inner_unnamed( client: &deadpool_postgres::Client, sql: &str, row_limit: Option, on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, ) -> Result { - match stream_select_query_prepared(client, sql, row_limit, on_item).await { + stream_select_query_inner_with_mode(client, sql, row_limit, on_item, true).await +} + +async fn stream_select_query_inner_with_mode( + client: &deadpool_postgres::Client, + sql: &str, + row_limit: Option, + on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, + force_unnamed: bool, +) -> Result { + match stream_select_query_prepared(client, sql, row_limit, on_item, force_unnamed).await { Ok(rows) => Ok(rows), Err(PostgresQueryStreamError::TextFallback { column_types, unsupported_type }) => { log::info!( @@ -1548,7 +1707,7 @@ pub(crate) async fn stream_select_query_inner( // Evict and retry once, matching the normal query execution path. log::warn!("[postgres][stream:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); client.statement_cache.remove(sql, &[]); - match stream_select_query_prepared(client, sql, row_limit, on_item).await { + match stream_select_query_prepared(client, sql, row_limit, on_item, force_unnamed).await { Ok(rows) => Ok(rows), Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_text_query(&err) => @@ -1596,17 +1755,20 @@ async fn stream_query_rows_on_client( cancelled: &AtomicBool, on_row: &mut impl FnMut(&[serde_json::Value]) -> Result<(), String>, ) -> Result { - let (stmt, metadata) = prepare_select_with_metadata(client, sql).await.map_err(pg_error_to_string)?; - let PreparedSelectMetadata { column_classes, unsupported_type, .. } = metadata; - if let Some(unsupported_type) = unsupported_type { - log::info!( - "[postgres][row_stream:text_fallback] unsupported_type={} switching_to=simple_query", - unsupported_type - ); - return stream_query_rows_text_on_client(client, sql, max_rows, cancelled, Some(&column_classes), on_row).await; - } - let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); - let stream = client.query_raw(&stmt, params).await.map_err(pg_error_to_string)?; + let stream_outcome = start_postgres_select_stream(client, sql, false).await.map_err(pg_error_to_string)?; + let (stream, metadata) = match stream_outcome { + PostgresSelectStreamOutcome::Binary { stream, metadata } => (stream, metadata), + PostgresSelectStreamOutcome::TextFallback { column_types, unsupported_type } => { + log::info!( + "[postgres][row_stream:text_fallback] unsupported_type={} switching_to=simple_query", + unsupported_type + ); + let column_classes = classify_pg_column_types(&column_types); + return stream_query_rows_text_on_client(client, sql, max_rows, cancelled, Some(&column_classes), on_row) + .await; + } + }; + let PreparedSelectMetadata { column_classes, .. } = metadata; tokio::pin!(stream); let row_limit = max_rows.unwrap_or(usize::MAX); let mut rows_exported = 0_u64; @@ -5488,7 +5650,7 @@ pub async fn execute_query_in_read_only_transaction_with_rollback( execute_postgres_infra_statement(&client, &statement, budget.recycle_timeout, stage).await?; } - execute_postgres_user_query( + execute_postgres_user_query_with_mode( &client, sql, max_rows, @@ -5497,6 +5659,7 @@ pub async fn execute_query_in_read_only_transaction_with_rollback( budget.cancel_timeout, cancel_context, false, + true, ) .await }, @@ -5571,7 +5734,13 @@ pub async fn stream_select_query_with_cancel( Ok(()) }; let result = await_stream_with_progress_timeout( - stream_select_query_inner(&client, sql, row_limit, &mut on_stream_item), + stream_select_query_inner_with_mode( + &client, + sql, + row_limit, + &mut on_stream_item, + setup_transaction_started, + ), query_timeout, progress_clock, cancel_token.as_ref(), @@ -5637,7 +5806,7 @@ pub async fn execute_query_with_schema_and_max_rows( "[postgres][execute_with_schema:skip-search-path] total_ms={} reason=transaction-recovery", start.elapsed().as_millis() ); - return execute_query_with_max_rows_inner(&client, sql, max_rows, false, None).await; + return execute_query_with_max_rows_inner(&client, sql, max_rows, false, None, false).await; } let set_schema_start = Instant::now(); @@ -5649,7 +5818,7 @@ pub async fn execute_query_with_schema_and_max_rows( ); let query_start = Instant::now(); - let result = execute_query_with_max_rows_inner(&client, sql, max_rows, false, None).await; + let result = execute_query_with_max_rows_inner(&client, sql, max_rows, false, None, false).await; if result.is_ok() { clear_postgres_caches_after_ddl(pool, Some(&client), sql); } @@ -5784,7 +5953,7 @@ pub(crate) async fn execute_postgres_infra_statement( timeout_duration: Duration, stage: &str, ) -> Result { - tokio::time::timeout(timeout_duration, client.execute(sql, &[])) + tokio::time::timeout(timeout_duration, client.execute_typed(sql, &[])) .await .map_err(|_| format!("PostgreSQL {stage} timed out after {} seconds", timeout_duration.as_secs()))? .map_err(pg_error_to_string) @@ -5861,6 +6030,31 @@ async fn execute_postgres_user_query( cancel_timeout: Duration, cancel_context: Option, prefer_text_protocol: bool, +) -> Result { + execute_postgres_user_query_with_mode( + client, + sql, + max_rows, + cancel_token, + timeout_duration, + cancel_timeout, + cancel_context, + prefer_text_protocol, + false, + ) + .await +} + +async fn execute_postgres_user_query_with_mode( + client: &deadpool_postgres::Client, + sql: &str, + max_rows: Option, + cancel_token: Option, + timeout_duration: Option, + cancel_timeout: Duration, + cancel_context: Option, + prefer_text_protocol: bool, + force_unnamed: bool, ) -> Result { let pg_cancel_token = client.cancel_token(); // Commands do not expose incremental results, so keep their timeout as a @@ -5874,7 +6068,7 @@ async fn execute_postgres_user_query( cancel_token, timeout_duration, cancel_timeout, - execute_query_with_max_rows_inner(client, sql, max_rows, prefer_text_protocol, None), + execute_query_with_max_rows_inner(client, sql, max_rows, prefer_text_protocol, None, force_unnamed), ) .await; } @@ -5883,7 +6077,14 @@ async fn execute_postgres_user_query( format!("Query timed out after {} seconds", timeout_duration.map_or(0, |timeout| timeout.as_secs())); let progress_clock = Arc::new(StreamProgressClock::new()); let result = await_stream_with_progress_timeout( - execute_query_with_max_rows_inner(client, sql, max_rows, prefer_text_protocol, Some(progress_clock.clone())), + execute_query_with_max_rows_inner( + client, + sql, + max_rows, + prefer_text_protocol, + Some(progress_clock.clone()), + force_unnamed, + ), timeout_duration, progress_clock, cancel_token.as_ref(), @@ -6062,6 +6263,7 @@ async fn execute_query_with_max_rows_inner( max_rows: Option, prefer_text_protocol: bool, progress_clock: Option>, + force_unnamed: bool, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); @@ -6075,10 +6277,13 @@ async fn execute_query_with_max_rows_inner( if prefer_text_protocol { execute_select_text(client, sql, start, row_limit, None, progress_clock.as_deref()).await } else { - execute_select_query_with_progress(client, sql, start, row_limit, progress_clock.as_deref()).await + execute_select_query_with_progress(client, sql, start, row_limit, progress_clock.as_deref(), force_unnamed) + .await } } else { - client.execute(sql, &[]).await.map_err(pg_error_to_string).map(|affected| QueryResult { + let affected = + if force_unnamed { client.execute_typed(sql, &[]).await } else { client.execute(sql, &[]).await }; + affected.map_err(pg_error_to_string).map(|affected| QueryResult { columns: vec![], column_types: Vec::new(), column_sortables: Vec::new(), @@ -6983,6 +7188,7 @@ mod tests { None, false, None, + false, ) .await .expect("execute statement with notice"); @@ -7785,12 +7991,18 @@ mod tests { let client = checkout_postgres_client(&pool, None, Duration::from_secs(5)).await?; let mut query_export_rows = Vec::new(); - stream_select_query_inner(&client, &select_sql, None, &mut |item| { - if let PostgresQueryStreamItem::Row(row) = item { - query_export_rows.push(row); - } - Ok(()) - }) + stream_select_query_inner_with_mode( + &client, + &select_sql, + None, + &mut |item| { + if let PostgresQueryStreamItem::Row(row) = item { + query_export_rows.push(row); + } + Ok(()) + }, + false, + ) .await?; let cancelled = AtomicBool::new(false); @@ -10670,6 +10882,37 @@ mod tests { )); } + #[test] + fn postgres_missing_prepared_statement_fallback_uses_sqlstate_and_message() { + assert!(should_fallback_postgres_missing_prepared_statement_fields(Some("26000"), "本地化的预编译语句错误",)); + assert!(should_fallback_postgres_missing_prepared_statement_fields( + None, + "prepared statement \"s63\" does not exist", + )); + } + + #[test] + fn postgres_missing_prepared_statement_fallback_rejects_unrelated_errors() { + assert!(!should_fallback_postgres_missing_prepared_statement_fields( + Some("23505"), + "duplicate key value violates unique constraint", + )); + assert!(!should_fallback_postgres_missing_prepared_statement_fields( + None, + "prepared statement result type changed", + )); + } + + #[test] + fn postgres_typed_params_pairs_values_with_inferred_types() { + let id = 42_i32; + let name = "dbx"; + let params: &[&(dyn tokio_postgres::types::ToSql + Sync)] = &[&id, &name]; + let typed = postgres_typed_params(params, &[Type::INT4, Type::TEXT]).expect("matching parameter count"); + assert_eq!(typed.iter().map(|(_, ty)| ty.clone()).collect::>(), vec![Type::INT4, Type::TEXT]); + assert!(postgres_typed_params(params, &[Type::INT4]).is_none()); + } + // --- execute_batch --- #[tokio::test] @@ -10763,6 +11006,120 @@ mod tests { assert_ne!(timezone, "Invalid/DBX_Timezone"); } + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] + async fn cached_queries_downgrade_connection_after_server_deallocates_statements() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") + .await + .expect("connect PostgreSQL database"); + let client = pool.get().await.expect("checkout postgres"); + let sql = "SELECT $1::int4 AS value"; + + let rows = postgres_query_cached(&client, sql, &[&41_i32]).await.expect("prime statement cache"); + assert_eq!(rows[0].get::<_, i32>(0), 41); + client.batch_execute("DEALLOCATE ALL").await.expect("drop server-side prepared statements"); + + let rows = postgres_query_cached(&client, sql, &[&42_i32]).await.expect("retry through unnamed typed query"); + assert_eq!(rows[0].get::<_, i32>(0), 42); + assert!(postgres_client_uses_unnamed_statements(&client)); + + client.execute_typed("DEALLOCATE ALL", &[]).await.expect("clear any remaining named statements"); + let rows = + postgres_query_cached(&client, sql, &[&43_i32]).await.expect("reuse inferred types without named bind"); + assert_eq!(rows[0].get::<_, i32>(0), 43); + + let prepared_count = client + .query_typed_one("SELECT count(*)::int8 FROM pg_prepared_statements", &[]) + .await + .expect("count server-side prepared statements") + .get::<_, i64>(0); + assert_eq!(prepared_count, 0, "downgraded connection must not bind or re-prepare the cached query"); + } + + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] + async fn missing_statement_fallback_uses_current_metadata_for_empty_results() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") + .await + .expect("connect PostgreSQL database"); + let client = pool.get().await.expect("checkout postgres"); + let select_sql = "SELECT * FROM dbx_statement_metadata_test WHERE false"; + + client + .execute_typed("CREATE TEMP TABLE dbx_statement_metadata_test (id int4)", &[]) + .await + .expect("create temporary test table"); + let initial = + execute_select_query(&client, select_sql, Instant::now(), 10).await.expect("prime statement cache"); + assert_eq!(initial.columns, vec!["id"]); + assert!(initial.rows.is_empty()); + + client + .execute_typed("ALTER TABLE dbx_statement_metadata_test ADD COLUMN label text", &[]) + .await + .expect("change result metadata"); + client.execute_typed("DEALLOCATE ALL", &[]).await.expect("drop server-side prepared statements"); + + let current = execute_select_query(&client, select_sql, Instant::now(), 10) + .await + .expect("fallback with current row description"); + assert_eq!(current.columns, vec!["id", "label"]); + assert_eq!(current.column_types, vec!["int4", "text"]); + assert!(current.rows.is_empty()); + } + + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] + async fn explicit_transaction_queries_start_unnamed_and_remain_usable() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") + .await + .expect("connect PostgreSQL database"); + let client = pool.get().await.expect("checkout postgres"); + + client.execute_typed("BEGIN", &[]).await.expect("begin transaction"); + client.execute_typed("DEALLOCATE ALL", &[]).await.expect("simulate statement loss"); + let result = execute_select_query_unnamed(&client, "SELECT 44::int4 AS value", Instant::now(), 10) + .await + .expect("execute unnamed query inside transaction"); + assert_eq!(result.rows[0][0], serde_json::json!(44)); + client.execute_typed("SELECT 1", &[]).await.expect("transaction remains usable"); + client.execute_typed("ROLLBACK", &[]).await.expect("rollback transaction"); + } + + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] + async fn backup_snapshot_stream_starts_unnamed_and_remains_usable() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") + .await + .expect("connect PostgreSQL database"); + let client = pool.get().await.expect("checkout postgres"); + + client + .execute_typed("BEGIN TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY", &[]) + .await + .expect("begin backup snapshot"); + client.execute_typed("DEALLOCATE ALL", &[]).await.expect("simulate statement loss"); + let mut columns = Vec::new(); + let mut rows = Vec::new(); + stream_select_query_inner_unnamed(&client, "SELECT 45::int4 AS value", None, &mut |item| { + match item { + PostgresQueryStreamItem::Columns { columns: names, .. } => columns = names, + PostgresQueryStreamItem::Row(row) => rows.push(row), + } + Ok(()) + }) + .await + .expect("stream unnamed query inside backup snapshot"); + assert_eq!(columns, vec!["value"]); + assert_eq!(rows[0][0], serde_json::json!(45)); + client.execute_typed("SELECT 1", &[]).await.expect("snapshot remains usable"); + client.execute_typed("ROLLBACK", &[]).await.expect("rollback backup snapshot"); + } + #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a PostgreSQL database"] async fn explicit_timezone_remains_strict_and_overrides_local_timezone() { diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index bfebb98b3d..100a1a3c3d 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4309,7 +4309,7 @@ async fn begin_transaction_session( TxnPoolHandle::Postgres(pg_pool) => { let conn = pg_pool.get().await.map_err(|e| format!("Failed to get Postgres connection: {e}"))?; let begin_sql = postgres_transaction_begin_sql(consistent_snapshot); - conn.execute(begin_sql, &[]).await.map_err(|e| format!("BEGIN failed: {e}"))?; + conn.execute_typed(begin_sql, &[]).await.map_err(|e| format!("BEGIN failed: {e}"))?; if let Some(schema) = schema { db::postgres::set_postgres_search_path( &conn, @@ -4595,7 +4595,7 @@ where TxnConnection::Postgres(conn) => { let mut batch = Vec::with_capacity(batch_size); let mut total_rows = 0_u64; - let result = db::postgres::stream_select_query_inner(conn, sql, None, &mut |item| { + let result = db::postgres::stream_select_query_inner_unnamed(conn, sql, None, &mut |item| { if let db::postgres::PostgresQueryStreamItem::Row(row) = item { batch.push(row); total_rows += 1; @@ -4682,7 +4682,7 @@ where async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), String> { match conn { TxnConnection::Postgres(conn) => { - conn.execute("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?; + conn.execute_typed("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?; } TxnConnection::Mysql(conn) => { conn.query_drop("ROLLBACK").await.map_err(|e| format!("ROLLBACK failed: {e}"))?; @@ -4740,9 +4740,9 @@ async fn execute_manual_txn_postgres_statement( row_limit: usize, ) -> Result { if db::postgres::postgres_statement_returns_rows(sql) { - db::postgres::execute_select_query(conn, sql, std::time::Instant::now(), row_limit).await + db::postgres::execute_select_query_unnamed(conn, sql, std::time::Instant::now(), row_limit).await } else { - let affected = conn.execute(sql, &[]).await.map_err(|e| format!("Query failed: {e}"))?; + let affected = conn.execute_typed(sql, &[]).await.map_err(|e| format!("Query failed: {e}"))?; Ok(db::QueryResult { columns: vec![], column_types: Vec::new(), @@ -4833,7 +4833,7 @@ pub async fn commit_manual_transaction(state: &AppState, txn_session_id: &str) - let mut conn = session.connection.lock().await; match &mut *conn { TxnConnection::Postgres(conn) => { - conn.execute("COMMIT", &[]).await.map_err(|e| format!("COMMIT failed: {e}"))?; + conn.execute_typed("COMMIT", &[]).await.map_err(|e| format!("COMMIT failed: {e}"))?; } TxnConnection::Mysql(conn) => { conn.query_drop("COMMIT").await.map_err(|e| format!("COMMIT failed: {e}"))?; diff --git a/crates/dbx-core/tests/live_manual_transaction.rs b/crates/dbx-core/tests/live_manual_transaction.rs index 0e382fdaf6..5ba20d6c27 100644 --- a/crates/dbx-core/tests/live_manual_transaction.rs +++ b/crates/dbx-core/tests/live_manual_transaction.rs @@ -3,6 +3,7 @@ use dbx_core::database_export::{begin_database_backup_snapshot_core, export_data use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; use dbx_core::query::{ begin_manual_transaction, commit_manual_transaction, execute_in_manual_transaction, rollback_manual_transaction, + stream_rows_in_manual_transaction, }; use dbx_core::storage::Storage; use std::time::{Duration, Instant}; @@ -16,6 +17,7 @@ fn live_config(prefix: &str, db_type: DatabaseType, default_port: u16) -> Connec let username = std::env::var(format!("{prefix}_USER")).expect("live DB user env var"); let password = std::env::var(format!("{prefix}_PASSWORD")).expect("live DB password env var"); let database = std::env::var(format!("{prefix}_DATABASE")).expect("live DB database env var"); + let url_params = std::env::var(format!("{prefix}_URL_PARAMS")).ok(); serde_json::from_value(serde_json::json!({ "id": format!("manual-txn-{prefix}"), @@ -29,7 +31,8 @@ fn live_config(prefix: &str, db_type: DatabaseType, default_port: u16) -> Connec "connect_timeout_secs": 5, "query_timeout_secs": 30, "idle_timeout_secs": 60, - "keepalive_interval_secs": 0 + "keepalive_interval_secs": 0, + "url_params": url_params })) .expect("live connection config should deserialize") } @@ -50,6 +53,9 @@ async fn live_manual_transaction_postgres_preserves_typed_selects_and_empty_meta let (state, db_path) = app_state_with_config(config.clone()).await; let txn = begin_manual_transaction(&state, &config.id, &database, None, None).await.expect("begin"); + execute_in_manual_transaction(&state, &txn, "DEALLOCATE ALL", &database, None, Some(10)) + .await + .expect("simulate lost prepared statements"); let typed = execute_in_manual_transaction( &state, &txn, @@ -80,6 +86,38 @@ async fn live_manual_transaction_postgres_preserves_typed_selects_and_empty_meta let _ = std::fs::remove_file(db_path); } +#[tokio::test] +#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL"] +async fn live_postgres_backup_snapshot_streams_after_server_deallocates_statements() { + let config = live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432); + let database = config.database.clone().expect("database"); + let (state, db_path) = app_state_with_config(config.clone()).await; + + let snapshot = begin_database_backup_snapshot_core(&state, &config.id, &database).await.expect("begin snapshot"); + execute_in_manual_transaction(&state, &snapshot.session_id, "DEALLOCATE ALL", &database, None, Some(10)) + .await + .expect("simulate lost prepared statements"); + + let mut batches = Vec::new(); + let row_count = stream_rows_in_manual_transaction( + &state, + &snapshot.session_id, + "SELECT 1::int4 AS value UNION ALL SELECT 2::int4", + 1, + |batch| { + batches.push(batch); + Ok(()) + }, + ) + .await + .expect("stream through backup snapshot"); + assert_eq!(row_count, 2); + assert_eq!(batches, vec![vec![vec![serde_json::json!(1)]], vec![vec![serde_json::json!(2)]]]); + + rollback_manual_transaction(&state, &snapshot.session_id).await.expect("rollback snapshot"); + let _ = std::fs::remove_file(db_path); +} + #[tokio::test] #[ignore = "requires DBX_LIVE_MANUAL_TXN_MYSQL_* env vars pointing at writable MySQL"] async fn live_manual_transaction_mysql_streams_with_row_limit() {