1 parent
2305f2de22
commit
9e135f7609
4 files changed
+618
-54
No files matched your search
@@ -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<T, F>(
|
||||
export_id: &str,
|
||||
cancel_token: &CancellationToken,
|
||||
operation: F,
|
||||
) -> Result<T, String>
|
||||
where
|
||||
F: Future<Output = Result<T, String>>,
|
||||
{
|
||||
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()));
|
||||
|
||||
@@ -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<usize>,
|
||||
on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>,
|
||||
cancel_token: Option<&CancellationToken>,
|
||||
budget: &DbOperationBudget,
|
||||
cancel_context: Option<&PostgresCancelContext>,
|
||||
) -> Result<u64, String> {
|
||||
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::<Result<(), String>>().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::<Result<(), String>>(),
|
||||
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());
|
||||
|
||||
@@ -4753,6 +4753,20 @@ pub async fn stream_rows_in_manual_transaction<F>(
|
||||
txn_session_id: &str,
|
||||
sql: &str,
|
||||
batch_size: usize,
|
||||
on_batch: F,
|
||||
) -> Result<u64, String>
|
||||
where
|
||||
F: FnMut(Vec<Vec<serde_json::Value>>) -> 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<F>(
|
||||
state: &AppState,
|
||||
txn_session_id: &str,
|
||||
sql: &str,
|
||||
batch_size: usize,
|
||||
cancel_token: Option<CancellationToken>,
|
||||
mut on_batch: F,
|
||||
) -> Result<u64, String>
|
||||
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<Duration>,
|
||||
) -> 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}"))?;
|
||||
|
||||
@@ -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() {
|
||||
|
||||
Reference in new issue
Block a user