diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index da87a451fe..4bd8981e87 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -5,6 +5,7 @@ use std::future::Future; use std::io::{BufWriter, Write}; use std::sync::RwLock; use std::time::Duration; +use tokio_util::sync::CancellationToken; use crate::connection::task_client_session_id; use crate::models::connection::DatabaseType; @@ -1486,6 +1487,46 @@ where } } +fn snapshot_batch_cancelled(db_type: &DatabaseType, export_id: &str) -> bool { + *db_type != DatabaseType::Postgres && is_export_cancelled_now(export_id) +} + +/// Unlike metadata operations, a database row stream owns a live server query +/// and transaction cleanup. Signal its cancellation token, then keep polling +/// the future until that cleanup finishes before reporting Export cancelled. +async fn await_export_stream_operation( + export_id: &str, + cancel_token: &CancellationToken, + operation: F, +) -> Result +where + F: Future>, +{ + tokio::pin!(operation); + loop { + if is_export_cancelled_now(export_id) { + cancel_token.cancel(); + if let Err(error) = operation.await { + if error.contains("Transaction cleanup failed") || error.contains("stream cleanup timed out") { + log::warn!("[database-export] cancellation cleanup was incomplete export_id={export_id}: {error}"); + } + } + return Err(EXPORT_CANCELLED_ERROR.to_string()); + } + tokio::select! { + biased; + result = &mut operation => { + return if is_export_cancelled_now(export_id) { + Err(EXPORT_CANCELLED_ERROR.to_string()) + } else { + result + }; + }, + _ = tokio::time::sleep(EXPORT_CANCEL_POLL_INTERVAL) => {} + } + } +} + pub async fn begin_database_backup_snapshot_core( state: &crate::connection::AppState, connection_id: &str, @@ -2508,39 +2549,48 @@ async fn export_database_sql_core_inner( if !col_names.is_empty() { if let Some(snapshot_session_id) = request.snapshot_session_id.as_deref() { let sql = database_export_select_sql(&col_names, &col_types, table_name, &request.schema, &db_type); - crate::query::stream_rows_in_manual_transaction( - state, - snapshot_session_id, - &sql, - batch_size, - |rows| { - if is_export_cancelled_now(&request.export_id) { - return Err("Export cancelled".to_string()); - } - write_database_export_rows( - &mut file, - &rows, - &col_names, - &col_types, - &col_extras, - table_name, - &request.schema, - &db_type, - )?; - total_rows_exported += rows.len() as u64; - on_progress(ExportProgress { - export_id: request.export_id.clone(), - current_object: table_name.clone(), - object_index, - total_objects, - rows_exported: total_rows_exported, - total_rows: None, - status: ExportStatus::Running, - error: None, - preparing: false, - }); - Ok(()) - }, + let cancel_token = CancellationToken::new(); + await_export_stream_operation( + &request.export_id, + &cancel_token, + crate::query::stream_rows_in_manual_transaction_with_cancel( + state, + snapshot_session_id, + &sql, + batch_size, + Some(cancel_token.clone()), + |rows| { + // PostgreSQL cancellation must unwind through its CancelRequest path so the + // connection reaches ReadyForQuery before rollback. Other snapshot streams do + // not consume the token yet, so retain their existing next-batch cancellation. + if snapshot_batch_cancelled(&db_type, &request.export_id) { + return Err(EXPORT_CANCELLED_ERROR.to_string()); + } + write_database_export_rows( + &mut file, + &rows, + &col_names, + &col_types, + &col_extras, + table_name, + &request.schema, + &db_type, + )?; + total_rows_exported += rows.len() as u64; + on_progress(ExportProgress { + export_id: request.export_id.clone(), + current_object: table_name.clone(), + object_index, + total_objects, + rows_exported: total_rows_exported, + total_rows: None, + status: ExportStatus::Running, + error: None, + preparing: false, + }); + Ok(()) + }, + ), ) .await?; } else { @@ -2907,9 +2957,10 @@ fn build_database_export_object_source_sql( #[cfg(test)] mod tests { use super::{ - await_export_operation, clear_export_cancelled, concurrent_metadata_prefetch_allowed, - database_export_metadata_prefetch_concurrency, emit_database_export_cancelled, set_export_cancelled, - ExportStatus, EXPORT_CANCELLED_ERROR, + await_export_operation, await_export_stream_operation, clear_export_cancelled, + concurrent_metadata_prefetch_allowed, database_export_metadata_prefetch_concurrency, + emit_database_export_cancelled, set_export_cancelled, snapshot_batch_cancelled, ExportStatus, + EXPORT_CANCELLED_ERROR, }; use super::{ build_database_export_object_source_sql, build_database_sql_export, build_export_insert_statements, @@ -2930,7 +2981,11 @@ mod tests { use crate::storage::Storage; use crate::types::{ObjectInfo, ObjectSourceKind, TableInfo}; use serde_json::{json, Value}; - use std::sync::{Arc, Mutex}; + use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, + }; + use tokio_util::sync::CancellationToken; #[tokio::test] async fn await_export_operation_drops_pending_metadata_after_cancel() { @@ -2963,6 +3018,74 @@ mod tests { assert_eq!(result, Ok(42)); } + #[tokio::test] + async fn await_export_stream_signals_cancel_and_waits_for_cleanup() { + let export_id = format!("cancel-stream-{}", uuid::Uuid::new_v4()); + clear_export_cancelled(&export_id).await; + let cancel_token = CancellationToken::new(); + let stream_cancel_token = cancel_token.clone(); + let cleaned = Arc::new(AtomicBool::new(false)); + let cleaned_after_cancel = cleaned.clone(); + let task_export_id = export_id.clone(); + let task = tokio::spawn(async move { + await_export_stream_operation(&task_export_id, &cancel_token, async move { + stream_cancel_token.cancelled().await; + cleaned_after_cancel.store(true, Ordering::SeqCst); + Err::<(), _>(crate::query::QUERY_CANCELED.to_string()) + }) + .await + }); + + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + set_export_cancelled(&export_id).await; + let result = tokio::time::timeout(std::time::Duration::from_secs(1), task) + .await + .expect("export stream cancellation should finish cleanup") + .expect("cancellation task should not panic"); + + assert_eq!(result, Err(EXPORT_CANCELLED_ERROR.to_string())); + assert!(cleaned.load(Ordering::SeqCst)); + clear_export_cancelled(&export_id).await; + } + + #[tokio::test] + async fn await_export_stream_prioritizes_existing_cancel_over_completion() { + let export_id = format!("cancel-ready-stream-{}", uuid::Uuid::new_v4()); + clear_export_cancelled(&export_id).await; + set_export_cancelled(&export_id).await; + let cancel_token = CancellationToken::new(); + + let result = await_export_stream_operation(&export_id, &cancel_token, async { Ok::<_, String>(42) }).await; + + assert_eq!(result, Err(EXPORT_CANCELLED_ERROR.to_string())); + assert!(cancel_token.is_cancelled()); + clear_export_cancelled(&export_id).await; + } + + #[tokio::test] + async fn await_export_stream_preserves_normal_completion() { + let export_id = format!("complete-stream-{}", uuid::Uuid::new_v4()); + clear_export_cancelled(&export_id).await; + let cancel_token = CancellationToken::new(); + + let result = await_export_stream_operation(&export_id, &cancel_token, async { Ok::<_, String>(42) }).await; + + assert_eq!(result, Ok(42)); + assert!(!cancel_token.is_cancelled()); + } + + #[tokio::test] + async fn non_postgres_snapshot_batches_keep_existing_cancel_check() { + let export_id = format!("cancel-non-postgres-batch-{}", uuid::Uuid::new_v4()); + clear_export_cancelled(&export_id).await; + set_export_cancelled(&export_id).await; + + assert!(snapshot_batch_cancelled(&DatabaseType::Mysql, &export_id)); + assert!(!snapshot_batch_cancelled(&DatabaseType::Postgres, &export_id)); + + clear_export_cancelled(&export_id).await; + } + #[test] fn cancelled_progress_is_terminal_and_has_no_error() { let events = Arc::new(Mutex::new(Vec::new())); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 26530c8d8a..2bf55a8b51 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1745,6 +1745,7 @@ async fn stream_select_query_text( Ok(rows_streamed) } +#[cfg(test)] pub(crate) async fn stream_select_query_inner_unnamed( client: &deadpool_postgres::Client, sql: &str, @@ -1754,6 +1755,57 @@ pub(crate) async fn stream_select_query_inner_unnamed( stream_select_query_inner_with_mode(client, sql, row_limit, on_item, true).await } +pub(crate) async fn stream_select_query_inner_unnamed_with_cancel( + client: &deadpool_postgres::Client, + sql: &str, + row_limit: Option, + on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, + cancel_token: Option<&CancellationToken>, + budget: &DbOperationBudget, + cancel_context: Option<&PostgresCancelContext>, +) -> Result { + if cancel_token.is_some_and(CancellationToken::is_cancelled) { + return Err(crate::query::canceled_error()); + } + let pg_cancel_token = client.cancel_token(); + let progress_clock = Arc::new(StreamProgressClock::new()); + let progress_clock_for_stream = progress_clock.clone(); + let mut on_stream_item = |item| { + let row_received = matches!(&item, PostgresQueryStreamItem::Row(_)); + on_item(item)?; + if row_received { + progress_clock_for_stream.mark(); + } + Ok(()) + }; + let timeout_error = + format!("Query timed out after {} seconds", budget.query_timeout.map_or(0, |timeout| timeout.as_secs())); + let stream = stream_select_query_inner_with_mode(client, sql, row_limit, &mut on_stream_item, true); + tokio::pin!(stream); + let result = await_stream_with_progress_timeout( + stream.as_mut(), + budget.query_timeout, + progress_clock, + cancel_token, + timeout_error.clone(), + ) + .await; + + if result.as_ref().is_err_and(|error| error == &timeout_error || error == crate::query::QUERY_CANCELED) { + let original_error = result.as_ref().expect_err("timeout or cancellation result").clone(); + cancel_postgres_query(pg_cancel_token, cancel_context, budget.cancel_timeout).await; + if tokio::time::timeout(budget.cleanup_timeout, stream.as_mut()).await.is_err() { + return Err(format!( + "{}; PostgreSQL stream cleanup timed out after {} seconds", + original_error, + budget.cleanup_timeout.as_secs() + )); + } + } + + result +} + async fn stream_select_query_inner_with_mode( client: &deadpool_postgres::Client, sql: &str, @@ -9792,6 +9844,44 @@ mod tests { assert_eq!(error, "query inactivity timeout"); } + #[tokio::test] + async fn postgres_row_query_timeout_covers_stall_after_last_row() { + let progress_clock = Arc::new(StreamProgressClock::new()); + let progress_clock_for_query = progress_clock.clone(); + let error = await_stream_with_progress_timeout( + async move { + progress_clock_for_query.mark(); + std::future::pending::>().await + }, + Some(Duration::from_millis(10)), + progress_clock, + None, + "query completion inactivity timeout".to_string(), + ) + .await + .expect_err("a stream stalled after its last row should time out"); + + assert_eq!(error, "query completion inactivity timeout"); + } + + #[tokio::test] + async fn postgres_row_query_without_timeout_still_honors_cancellation() { + let progress_clock = Arc::new(StreamProgressClock::new()); + let cancel_token = CancellationToken::new(); + cancel_token.cancel(); + let error = await_stream_with_progress_timeout( + std::future::pending::>(), + None, + progress_clock, + Some(&cancel_token), + "disabled query timeout".to_string(), + ) + .await + .expect_err("cancellation must remain active when query timeout is disabled"); + + assert_eq!(error, crate::query::QUERY_CANCELED); + } + #[test] fn postgres_cancel_context_omits_disabled_ssl_mode() { assert!(build_postgres_cancel_context("postgres://localhost/app?sslmode=disable").is_none()); diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index efe051972d..b23b6c8292 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -4753,6 +4753,20 @@ pub async fn stream_rows_in_manual_transaction( txn_session_id: &str, sql: &str, batch_size: usize, + on_batch: F, +) -> Result +where + F: FnMut(Vec>) -> Result<(), String> + Send, +{ + stream_rows_in_manual_transaction_with_cancel(state, txn_session_id, sql, batch_size, None, on_batch).await +} + +pub(crate) async fn stream_rows_in_manual_transaction_with_cancel( + state: &AppState, + txn_session_id: &str, + sql: &str, + batch_size: usize, + cancel_token: Option, mut on_batch: F, ) -> Result where @@ -4782,30 +4796,46 @@ where return Err(MANUAL_TRANSACTION_IDLE_ROLLBACK_ERROR.to_string()); } - let connection = { + let (connection, pool_key, connection_id) = { let sessions = state.transaction_sessions.read().await; sessions .get(txn_session_id) - .map(|session| Arc::clone(&session.connection)) + .map(|session| (Arc::clone(&session.connection), session.pool_key.clone(), session.connection_id.clone())) .ok_or(MANUAL_TRANSACTION_SESSION_NOT_FOUND_ERROR)? }; + let operation_budget = { + let configs = state.configs.read().await; + configs + .get(&connection_id) + .map(DbOperationBudget::from_connection_config) + .unwrap_or_else(DbOperationBudget::with_defaults) + }; + let postgres_cancel_context = state.get_postgres_cancel_context(&pool_key).await; let batch_size = batch_size.max(1); let mut conn = connection.lock().await; let stream_result = match &mut *conn { 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_unnamed(conn, sql, None, &mut |item| { - if let db::postgres::PostgresQueryStreamItem::Row(row) = item { - batch.push(row); - total_rows += 1; - if batch.len() >= batch_size { - on_batch(std::mem::take(&mut batch))?; - batch = Vec::with_capacity(batch_size); + let result = db::postgres::stream_select_query_inner_unnamed_with_cancel( + conn, + sql, + None, + &mut |item| { + if let db::postgres::PostgresQueryStreamItem::Row(row) = item { + batch.push(row); + total_rows += 1; + if batch.len() >= batch_size { + on_batch(std::mem::take(&mut batch))?; + batch = Vec::with_capacity(batch_size); + } } - } - Ok(()) - }) + Ok(()) + }, + cancel_token.as_ref(), + &operation_budget, + postgres_cancel_context.as_ref(), + ) .await; match result { Ok(_) if !batch.is_empty() => on_batch(batch).map(|_| total_rows), @@ -4866,8 +4896,13 @@ where sessions.remove(txn_session_id) }; if let Some(session) = removed { - let _ = rollback_manual_txn_connection(&mut conn).await; + let rollback_result = + rollback_manual_txn_connection_with_postgres_timeout(&mut conn, Some(operation_budget.cleanup_timeout)) + .await; release_manual_txn_session_pool(state, &session.connection_id, &mut conn).await; + if let Err(rollback_error) = rollback_result { + return Err(format!("{err}. Transaction cleanup failed: {rollback_error}")); + } } return Err(format!("{err}. Transaction was auto-rolled back.")); } @@ -4890,9 +4925,22 @@ where } async fn rollback_manual_txn_connection(conn: &mut TxnConnection) -> Result<(), String> { + rollback_manual_txn_connection_with_postgres_timeout(conn, None).await +} + +async fn rollback_manual_txn_connection_with_postgres_timeout( + conn: &mut TxnConnection, + postgres_timeout: Option, +) -> Result<(), String> { match conn { TxnConnection::Postgres(conn) => { - conn.execute_typed("ROLLBACK", &[]).await.map_err(|e| format!("ROLLBACK failed: {e}"))?; + if let Some(timeout) = postgres_timeout { + db::postgres::execute_postgres_infra_statement(conn, "ROLLBACK", timeout, "manual_txn.rollback") + .await + .map_err(|error| format!("ROLLBACK failed: {error}"))?; + } else { + 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}"))?; diff --git a/crates/dbx-core/tests/live_manual_transaction.rs b/crates/dbx-core/tests/live_manual_transaction.rs index 5ba20d6c27..c2c88e0d50 100644 --- a/crates/dbx-core/tests/live_manual_transaction.rs +++ b/crates/dbx-core/tests/live_manual_transaction.rs @@ -1,11 +1,18 @@ use dbx_core::connection::AppState; -use dbx_core::database_export::{begin_database_backup_snapshot_core, export_database_sql_core, DatabaseExportRequest}; +use dbx_core::database_export::{ + begin_database_backup_snapshot_core, clear_export_cancelled, export_database_sql_core, set_export_cancelled, + DatabaseExportRequest, ExportStatus, +}; 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, + begin_manual_transaction, commit_manual_transaction, execute_in_manual_transaction, execute_sql_statement, + rollback_manual_transaction, stream_rows_in_manual_transaction, }; use dbx_core::storage::Storage; +use std::sync::{ + atomic::{AtomicU64, Ordering}, + Arc, +}; use std::time::{Duration, Instant}; fn live_config(prefix: &str, db_type: DatabaseType, default_port: u16) -> ConnectionConfig { @@ -114,10 +121,306 @@ async fn live_postgres_backup_snapshot_streams_after_server_deallocates_statemen assert_eq!(row_count, 2); assert_eq!(batches, vec![vec![vec![serde_json::json!(1)]], vec![vec![serde_json::json!(2)]]]); + let empty_rows = stream_rows_in_manual_transaction( + &state, + &snapshot.session_id, + "SELECT 1::int4 AS value WHERE false", + 1, + |_| panic!("empty result must not emit a row batch"), + ) + .await + .expect("stream empty result through backup snapshot"); + assert_eq!(empty_rows, 0); + + let mut fallback_batches = Vec::new(); + let fallback_rows = stream_rows_in_manual_transaction( + &state, + &snapshot.session_id, + "SELECT ROW(1, 'fallback') AS value", + 1, + |batch| { + fallback_batches.push(batch); + Ok(()) + }, + ) + .await + .expect("stream text-fallback value through backup snapshot"); + assert_eq!(fallback_rows, 1); + assert_eq!(fallback_batches.len(), 1); + assert!(fallback_batches[0][0][0].as_str().is_some()); + 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_POSTGRES_* env vars pointing at writable PostgreSQL"] +async fn live_postgres_backup_snapshot_times_out_when_row_stream_stalls() { + let mut config = live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432); + config.query_timeout_secs = 1; + 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"); + let started_at = Instant::now(); + let result = stream_rows_in_manual_transaction( + &state, + &snapshot.session_id, + "SELECT pg_sleep(5), 1::int4 AS value", + 1, + |_| Ok(()), + ) + .await; + + assert_eq!(result, Err("Query timed out after 1 seconds. Transaction was auto-rolled back.".to_string())); + assert!(started_at.elapsed() < Duration::from_secs(5), "stalled snapshot stream was not cancelled promptly"); + assert!(!state.transaction_sessions.read().await.contains_key(&snapshot.session_id)); + + let recovery = begin_manual_transaction(&state, &config.id, &database, None, None).await.expect("begin recovery"); + execute_in_manual_transaction(&state, &recovery, "SELECT 1", &database, None, Some(1)) + .await + .expect("PostgreSQL should accept a query after snapshot timeout"); + rollback_manual_transaction(&state, &recovery).await.expect("rollback recovery"); + let _ = std::fs::remove_file(db_path); +} + +#[test] +#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL"] +fn live_postgres_backup_snapshot_exports_wide_jsonb_then_next_table() { + std::thread::Builder::new() + .name("dbx-live-pg-wide-backup".to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build() + .expect("build wide backup test runtime") + .block_on(live_postgres_backup_snapshot_exports_wide_jsonb_then_next_table_inner()) + }) + .expect("spawn wide backup test thread") + .join() + .expect("wide backup test thread should not panic"); +} + +async fn live_postgres_backup_snapshot_exports_wide_jsonb_then_next_table_inner() { + let mut config = live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432); + config.query_timeout_secs = 1; + let database = config.database.clone().expect("database"); + let connection_id = config.id.clone(); + let (state, db_path) = app_state_with_config(config).await; + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let schema = format!("dbx_7087_{}", &suffix[..12]); + let export_path = std::env::temp_dir().join(format!("dbx-7087-wide-jsonb-{suffix}.sql")); + + let setup_sql = [ + format!(r#"CREATE SCHEMA "{schema}""#), + format!(r#"CREATE TABLE "{schema}".audit_ops (id int PRIMARY KEY, task_payload jsonb NOT NULL)"#), + format!( + r#"INSERT INTO "{schema}".audit_ops + SELECT i, jsonb_build_object('id', i, 'contents', jsonb_build_array(repeat('x', 1500))) + FROM generate_series(1, 3000) AS source(i)"# + ), + format!(r#"CREATE TABLE "{schema}".zz_next (id int PRIMARY KEY, label text NOT NULL)"#), + format!( + r#"INSERT INTO "{schema}".zz_next + SELECT i, 'next-' || i FROM generate_series(1, 6) AS source(i)"# + ), + ]; + for sql in setup_sql { + execute_sql_statement(&state, &connection_id, &database, &sql, None, None) + .await + .expect("create wide JSONB backup fixture"); + } + + let snapshot = + begin_database_backup_snapshot_core(&state, &connection_id, &database).await.expect("begin snapshot"); + let request = DatabaseExportRequest { + export_id: format!("live-postgres-wide-jsonb-{suffix}"), + connection_id: connection_id.clone(), + database: database.clone(), + schema: schema.clone(), + file_path: export_path.to_string_lossy().to_string(), + selected_tables: vec!["audit_ops".to_string(), "zz_next".to_string()], + excluded_tables: Vec::new(), + include_structure: true, + include_data: true, + include_objects: false, + include_create_database: false, + drop_table_if_exists: false, + omit_auto_increment: false, + fail_on_error: true, + snapshot_session_id: Some(snapshot.session_id.clone()), + batch_size: 1000, + }; + let terminal_rows = AtomicU64::new(u64::MAX); + let export_result = export_database_sql_core(&state, &request, |progress| { + if matches!(progress.status, dbx_core::database_export::ExportStatus::Done) { + terminal_rows.store(progress.rows_exported, Ordering::Relaxed); + } + }) + .await; + let rollback_result = rollback_manual_transaction(&state, &snapshot.session_id).await; + + export_result.expect("export wide JSONB table and continue to next table"); + rollback_result.expect("rollback completed backup snapshot"); + assert_eq!(terminal_rows.load(Ordering::Relaxed), 3006); + let sql = std::fs::read_to_string(&export_path).expect("read wide JSONB export"); + assert!(sql.len() > 3_000_000, "wide JSONB fixture should produce a multi-megabyte export"); + assert!(sql.contains("audit_ops")); + assert!(sql.contains("zz_next")); + assert!(sql.contains("next-6"), "export must reach and finish the table after audit_ops"); + + execute_sql_statement(&state, &connection_id, &database, &format!(r#"DROP SCHEMA "{schema}" CASCADE"#), None, None) + .await + .expect("drop wide JSONB backup fixture"); + let _ = std::fs::remove_file(export_path); + let _ = std::fs::remove_file(db_path); +} + +#[test] +#[ignore = "requires DBX_LIVE_MANUAL_TXN_POSTGRES_* env vars pointing at writable PostgreSQL"] +fn live_postgres_backup_snapshot_cancel_interrupts_pending_row() { + std::thread::Builder::new() + .name("dbx-live-pg-backup-cancel".to_string()) + .stack_size(16 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .thread_stack_size(16 * 1024 * 1024) + .enable_all() + .build() + .expect("build backup cancellation test runtime") + .block_on(live_postgres_backup_snapshot_cancel_interrupts_pending_row_inner()) + }) + .expect("spawn backup cancellation test thread") + .join() + .expect("backup cancellation test thread should not panic"); +} + +async fn live_postgres_backup_snapshot_cancel_interrupts_pending_row_inner() { + let admin_config = live_config("DBX_LIVE_MANUAL_TXN_POSTGRES", DatabaseType::Postgres, 5432); + let database = admin_config.database.clone().expect("database"); + let admin_connection_id = admin_config.id.clone(); + let (admin_state, admin_db_path) = app_state_with_config(admin_config.clone()).await; + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let schema = format!("dbx_7087_cancel_{}", &suffix[..8]); + let role = format!("dbx7087_{}", &suffix[..12]); + let password = format!("dbx-{suffix}"); + + let setup_sql = [ + format!(r#"CREATE ROLE "{role}" LOGIN PASSWORD '{password}'"#), + format!(r#"CREATE SCHEMA "{schema}""#), + format!(r#"CREATE TABLE "{schema}".slow_rows (id int PRIMARY KEY, payload text NOT NULL)"#), + format!(r#"INSERT INTO "{schema}".slow_rows VALUES (1, repeat('x', 1000000)), (2, repeat('y', 1000000))"#), + format!( + r#"CREATE FUNCTION "{schema}".slow_visible(integer) RETURNS boolean + LANGUAGE plpgsql VOLATILE + AS 'BEGIN PERFORM pg_sleep(2); RETURN true; END'"# + ), + format!(r#"ALTER TABLE "{schema}".slow_rows ENABLE ROW LEVEL SECURITY"#), + format!( + r#"CREATE POLICY slow_rows_policy ON "{schema}".slow_rows + FOR SELECT TO "{role}" USING ("{schema}".slow_visible(id))"# + ), + format!(r#"GRANT USAGE ON SCHEMA "{schema}" TO "{role}""#), + format!(r#"GRANT SELECT ON "{schema}".slow_rows TO "{role}""#), + ]; + for sql in setup_sql { + execute_sql_statement(&admin_state, &admin_connection_id, &database, &sql, None, None) + .await + .expect("create cancellable backup fixture"); + } + + let mut role_config = admin_config; + role_config.id = format!("live-postgres-backup-cancel-{suffix}"); + role_config.name = "Live PostgreSQL backup cancellation".to_string(); + role_config.username = role.clone(); + role_config.password = password; + role_config.query_timeout_secs = 0; + let role_connection_id = role_config.id.clone(); + let (role_state, role_db_path) = app_state_with_config(role_config).await; + let role_state = Arc::new(role_state); + let snapshot = begin_database_backup_snapshot_core(&role_state, &role_connection_id, &database) + .await + .expect("begin role snapshot"); + let export_id = format!("live-postgres-backup-cancel-{suffix}"); + clear_export_cancelled(&export_id).await; + let export_path = std::env::temp_dir().join(format!("dbx-7087-cancel-{suffix}.sql")); + let request = DatabaseExportRequest { + export_id: export_id.clone(), + connection_id: role_connection_id, + database: database.clone(), + schema: schema.clone(), + file_path: export_path.to_string_lossy().to_string(), + selected_tables: vec!["slow_rows".to_string()], + excluded_tables: Vec::new(), + include_structure: false, + include_data: true, + include_objects: false, + include_create_database: false, + drop_table_if_exists: false, + omit_auto_increment: false, + fail_on_error: true, + snapshot_session_id: Some(snapshot.session_id.clone()), + batch_size: 1, + }; + let rows_seen = Arc::new(AtomicU64::new(0)); + let terminal_status = Arc::new(AtomicU64::new(0)); + let rows_for_export = rows_seen.clone(); + let terminal_for_export = terminal_status.clone(); + let state_for_export = role_state.clone(); + let export_task = tokio::spawn(async move { + export_database_sql_core(&state_for_export, &request, |progress| { + rows_for_export.store(progress.rows_exported, Ordering::Relaxed); + if matches!(progress.status, ExportStatus::Cancelled) { + terminal_for_export.store(1, Ordering::Relaxed); + } else if matches!(progress.status, ExportStatus::Done | ExportStatus::Error) { + terminal_for_export.store(2, Ordering::Relaxed); + } + }) + .await + }); + + tokio::time::timeout(Duration::from_secs(10), async { + while rows_seen.load(Ordering::Relaxed) == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("first slow row should be exported before cancellation"); + let cancel_started = Instant::now(); + set_export_cancelled(&export_id).await; + let export_result = tokio::time::timeout(Duration::from_secs(10), export_task) + .await + .expect("cancelled backup should terminate") + .expect("backup task should not panic"); + + assert_eq!(export_result, Ok(())); + assert_eq!(terminal_status.load(Ordering::Relaxed), 1, "export cancel must produce the Cancelled terminal state"); + assert!(cancel_started.elapsed() < Duration::from_secs(2), "CancelRequest did not interrupt the pending row"); + assert!(!role_state.transaction_sessions.read().await.contains_key(&snapshot.session_id)); + clear_export_cancelled(&export_id).await; + + execute_sql_statement( + &admin_state, + &admin_connection_id, + &database, + &format!(r#"DROP SCHEMA "{schema}" CASCADE"#), + None, + None, + ) + .await + .expect("drop cancellable backup schema"); + execute_sql_statement(&admin_state, &admin_connection_id, &database, &format!(r#"DROP ROLE "{role}""#), None, None) + .await + .expect("drop cancellable backup role"); + let _ = std::fs::remove_file(export_path); + let _ = std::fs::remove_file(role_db_path); + let _ = std::fs::remove_file(admin_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() {