diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index be3e9e6bb4..a04e167a08 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -4427,6 +4427,7 @@ async fn execute_result_set_with_text_protocol_on_conn( result_key_columns: &[String], diagnostic_trace_id: Option<&str>, start: Instant, + progress_clock: Option<&crate::query::StreamProgressClock>, ) -> Result { let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); let dispatch_started_at = diagnostics_enabled.then(Instant::now); @@ -4517,6 +4518,9 @@ async fn execute_result_set_with_text_protocol_on_conn( } let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; + if let Some(progress_clock) = progress_clock { + progress_clock.mark(); + } if result_rows.len() >= row_limit { truncated = true; break; @@ -4816,6 +4820,7 @@ async fn execute_result_set_with_prepared_protocol_on_conn( result_key_columns: &[String], diagnostic_trace_id: Option<&str>, start: Instant, + progress_clock: Option<&crate::query::StreamProgressClock>, ) -> Result { let diagnostics_enabled = diagnostic_trace_id.is_some() && log::log_enabled!(log::Level::Info); let dispatch_started_at = diagnostics_enabled.then(Instant::now); @@ -4848,6 +4853,9 @@ async fn execute_result_set_with_prepared_protocol_on_conn( } let Some(row) = next_row else { break }; let row = row.map_err(|e| e.to_string())?; + if let Some(progress_clock) = progress_clock { + progress_clock.mark(); + } if result_rows.len() >= row_limit { truncated = true; break; @@ -4958,6 +4966,54 @@ where execute_query_on_conn_with_max_rows(&mut conn, sql, bare, max_rows, dialect).await } +/// Progress-aware variant of [`execute_query_with_max_rows`] for long transfers. +/// +/// The statement runs under an inactivity budget: the clock is reset for every +/// row MySQL delivers, so a large table that keeps streaming is never cancelled +/// merely for exceeding the timeout in total — only a genuine stall is. +/// +/// Unlike the Postgres/SQL Server variants there is no returns-rows gate here: +/// a write delivers no rows, so no progress mark ever resets the clock and the +/// budget degrades to the same plain wall-clock timeout anyway. +pub(crate) async fn execute_query_with_max_rows_progress

( + pool: &P, + sql: &str, + bare: bool, + max_rows: Option, + dialect: MySqlQueryDialect, + progress_clock: std::sync::Arc, + timeout: Option, +) -> Result +where + P: MySqlPoolAccess + ?Sized, +{ + let mut conn = get_conn_with_health_check(pool).await?; + let timeout_error = format!("Query timed out after {} seconds", timeout.map_or(0, |timeout| timeout.as_secs())); + let clock_for_query = progress_clock.clone(); + crate::query::await_stream_with_progress_timeout( + async move { + execute_query_on_conn_with_limits_progress( + &mut conn, + sql, + bare, + max_rows, + None, + &[], + dialect, + None, + Some(&clock_for_query), + ) + .await + .map(|result| result.result) + }, + timeout, + progress_clock, + None, + timeout_error, + ) + .await +} + pub async fn stream_query_rows( pool: &MySqlPool, sql: &str, @@ -5163,6 +5219,37 @@ pub async fn execute_query_on_conn_with_limits( result_key_columns: &[String], dialect: MySqlQueryDialect, diagnostic_trace_id: Option<&str>, +) -> Result { + execute_query_on_conn_with_limits_progress( + conn, + sql, + bare, + max_rows, + max_result_bytes, + result_key_columns, + dialect, + diagnostic_trace_id, + None, + ) + .await +} + +/// [`execute_query_on_conn_with_limits`] with an optional progress clock. +/// +/// The clock is marked for every row the server delivers, so a caller wrapping +/// this in an inactivity budget (see [`crate::query::await_stream_with_progress_timeout`]) +/// only times out on a genuine stall, not on a long-but-steady stream — which is +/// what lets a transfer of a large table survive the configured timeout. +pub(crate) async fn execute_query_on_conn_with_limits_progress( + conn: &mut mysql_async::Conn, + sql: &str, + bare: bool, + max_rows: Option, + max_result_bytes: Option, + result_key_columns: &[String], + dialect: MySqlQueryDialect, + diagnostic_trace_id: Option<&str>, + progress_clock: Option<&crate::query::StreamProgressClock>, ) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); @@ -5178,6 +5265,7 @@ pub async fn execute_query_on_conn_with_limits( result_key_columns, diagnostic_trace_id, start, + progress_clock, ) .await } else { @@ -5189,6 +5277,7 @@ pub async fn execute_query_on_conn_with_limits( result_key_columns, diagnostic_trace_id, start, + progress_clock, ) .await { @@ -5203,6 +5292,7 @@ pub async fn execute_query_on_conn_with_limits( result_key_columns, diagnostic_trace_id, start, + progress_clock, ) .await } @@ -5922,6 +6012,53 @@ mod tests { new_row(values, columns.into()) } + #[tokio::test] + #[ignore = "requires a disposable MySQL reachable via DBX_LIVE_MYSQL_TRANSFER_{HOST,PORT,USER,PASSWORD}"] + async fn live_mysql_progress_read_survives_a_total_duration_beyond_the_timeout() { + let Ok(host) = std::env::var("DBX_LIVE_MYSQL_TRANSFER_HOST") else { + return; + }; + let port = std::env::var("DBX_LIVE_MYSQL_TRANSFER_PORT").unwrap_or_else(|_| "3306".to_string()); + let user = std::env::var("DBX_LIVE_MYSQL_TRANSFER_USER").unwrap_or_else(|_| "root".to_string()); + let password = std::env::var("DBX_LIVE_MYSQL_TRANSFER_PASSWORD").unwrap_or_default(); + let database = std::env::var("DBX_LIVE_MYSQL_TRANSFER_DATABASE").unwrap_or_else(|_| "test".to_string()); + let url = format!("mysql://{user}:{password}@{host}:{port}/{database}"); + let pool = connect(&url, Duration::from_secs(5)).await.unwrap(); + + // 20 rows produced ~50 ms apart, each carrying >8 KB so the server flushes + // it per row instead of buffering the small result set. The SLEEP sits in + // the WHERE clause so MySQL evaluates it per row — in the SELECT list the + // planner folds the constant call and the whole statement returns at once. + let sql = "WITH RECURSIVE seq(n) AS (SELECT 1 UNION ALL SELECT n + 1 FROM seq WHERE n < 20) \ + SELECT REPEAT('x', 20000) AS payload, n FROM seq WHERE SLEEP(0.05) = 0"; + + let progress_clock = std::sync::Arc::new(crate::query::StreamProgressClock::new()); + let result = execute_query_with_max_rows_progress( + &pool, + sql, + false, + None, + Default::default(), + progress_clock, + Some(Duration::from_millis(200)), + ) + .await; + assert!(result.is_ok(), "progress-aware read must survive a total duration beyond the timeout: {result:?}"); + assert_eq!(result.unwrap().rows.len(), 20); + + // Contrast: the same statement under a never-reset wall-clock budget must + // time out, proving this test actually exercises the difference. + let wall_clock = crate::query::await_stream_with_progress_timeout( + execute_query_with_max_rows(&pool, sql, false, None, Default::default()), + Some(Duration::from_millis(200)), + std::sync::Arc::new(crate::query::StreamProgressClock::new()), + None, + "Query timed out after 0 seconds".to_string(), + ) + .await; + assert!(wall_clock.is_err(), "the wall-clock path must still time out"); + } + #[test] fn mysql_first_column_value_ignores_extra_result_columns() { let row = mysql_async::from_row::(mysql_test_row(vec![ diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 1bd3574f8e..ad6d1c9137 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -6867,6 +6867,44 @@ pub async fn execute_query_with_max_rows( } } +/// Progress-aware variant of [`execute_query_with_max_rows`] for long transfers. +/// +/// Row-returning queries run under an *inactivity* budget: the clock is reset +/// every time PostgreSQL actually delivers a row, so a large table that keeps +/// streaming is never cancelled merely for exceeding the timeout in total. Only +/// a genuine stall (no row for the whole timeout) is reported as a timeout. +/// Statements that return no rows have no incremental progress to report, so +/// they keep the plain wall-clock path. +pub(crate) async fn execute_query_with_max_rows_progress( + pool: &Pool, + sql: &str, + max_rows: Option, + progress_clock: Arc, + timeout: Option, +) -> Result { + if !postgres_statement_returns_rows(sql) { + // DDL/DML expose no incremental progress, so keep the original wall-clock + // budget: a hung write must still be bounded by the configured timeout. + return crate::query::wait_for_query_opt(None, timeout, execute_query_with_max_rows(pool, sql, max_rows)).await; + } + + let start = Instant::now(); + let row_limit = query_result_row_limit(max_rows); + let timeout_error = format!("Query timed out after {} seconds", timeout.map_or(0, |timeout| timeout.as_secs())); + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let clock_for_select = progress_clock.clone(); + await_stream_with_progress_timeout( + async move { + execute_select_query_with_progress(&client, sql, start, row_limit, Some(&clock_for_select), false).await + }, + timeout, + progress_clock, + None, + timeout_error, + ) + .await +} + pub async fn execute_query_with_max_rows_and_cancel( pool: &Pool, sql: &str, @@ -11405,6 +11443,42 @@ mod tests { assert_eq!(result, Ok(())); } + #[tokio::test] + #[ignore = "requires DBX_LIVE_PG_TRANSFER_SOURCE_URL pointing at a disposable PostgreSQL"] + async fn live_postgres_progress_read_survives_a_total_duration_beyond_the_timeout() { + let Ok(url) = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL") else { + return; + }; + let pool = connect(&url, Duration::from_secs(5)).await.unwrap(); + + // 20 rows produced ~50 ms apart: the statement streams for ~1 s in total, + // well beyond the 200 ms budget, while never stalling that long between + // rows — the exact shape a progress-aware transfer read has to survive. + // Each row carries >8 KB so PostgreSQL flushes it immediately instead of + // buffering the whole (tiny) result set and sending it in one packet. + let sql = "SELECT pg_sleep(0.05) IS NULL AS slept, repeat('x', 20000) AS payload, n \ + FROM generate_series(1, 20) AS n"; + + let progress_clock = Arc::new(StreamProgressClock::new()); + let result = + execute_query_with_max_rows_progress(&pool, sql, None, progress_clock, Some(Duration::from_millis(200))) + .await; + assert!(result.is_ok(), "progress-aware read must survive a total duration beyond the timeout: {result:?}"); + assert_eq!(result.unwrap().rows.len(), 20); + + // Contrast: the same statement under a plain, never-reset wall-clock budget + // must time out, proving this test actually exercises the difference. + let wall_clock = await_stream_with_progress_timeout( + execute_query_with_max_rows(&pool, sql, None), + Some(Duration::from_millis(200)), + Arc::new(StreamProgressClock::new()), + None, + "Query timed out after 0 seconds".to_string(), + ) + .await; + assert!(wall_clock.is_err(), "the wall-clock path must still time out"); + } + #[tokio::test] async fn postgres_row_query_timeout_still_fires_after_no_progress() { let progress_clock = Arc::new(StreamProgressClock::new()); diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 66aa5eca27..7533b7a559 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -677,6 +677,31 @@ mod tests { assert_eq!(result.rows[0][0], serde_json::json!("Ada")); } + #[tokio::test] + async fn progress_read_marks_the_inactivity_clock_for_a_multi_row_select() { + let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); + execute_query(&pool, "CREATE TABLE progress_probe (id INTEGER PRIMARY KEY)").await.expect("create table"); + for id in 1..=5 { + execute_query(&pool, &format!("INSERT INTO progress_probe (id) VALUES ({id})")).await.expect("insert row"); + } + + let progress_clock = std::sync::Arc::new(crate::query::StreamProgressClock::new()); + assert!(!progress_clock.marked(), "a fresh clock must not report progress"); + + let result = execute_query_with_max_rows_progress( + &pool, + "SELECT id FROM progress_probe ORDER BY id", + None, + progress_clock.clone(), + Some(std::time::Duration::from_secs(30)), + ) + .await + .expect("progress read"); + + assert_eq!(result.rows.len(), 5); + assert!(progress_clock.marked(), "streamed rows must record progress so the inactivity budget resets"); + } + #[tokio::test] async fn sqlite_dml_returning_preserves_result_rows() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); @@ -2830,12 +2855,53 @@ pub async fn execute_query_with_max_rows( return worker.query(&sql, max_rows).await; } let pool = pool.clone(); - tokio::task::spawn_blocking(move || execute_query_blocking(&pool, &sql, max_rows)) + tokio::task::spawn_blocking(move || execute_query_blocking(&pool, &sql, max_rows, None)) .await .map_err(|e| e.to_string())? } -fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option) -> Result { +/// Progress-aware variant of [`execute_query_with_max_rows`] for long transfers. +/// +/// The clock is marked for every row the statement produces, so a caller wrapping +/// this in an inactivity budget only times out on a genuine stall, not on a long +/// but steady read. The SQLite worker (sidecar) path keeps the plain call: its RPC +/// returns the whole result at once and exposes no incremental progress. +pub(crate) async fn execute_query_with_max_rows_progress( + pool: &SqliteHandle, + sql: &str, + max_rows: Option, + progress_clock: std::sync::Arc, + timeout: Option, +) -> Result { + let sql = normalize_sqlite_sql(sql); + if let Some(worker) = pool.worker() { + // The sidecar worker returns the whole result in one RPC and exposes no + // incremental progress, so keep the original wall-clock budget. + return crate::query::wait_for_query_opt(None, timeout, worker.query(&sql, max_rows)).await; + } + let pool = pool.clone(); + let timeout_error = format!("Query timed out after {} seconds", timeout.map_or(0, |timeout| timeout.as_secs())); + let clock_for_query = progress_clock.clone(); + crate::query::await_stream_with_progress_timeout( + async move { + tokio::task::spawn_blocking(move || execute_query_blocking(&pool, &sql, max_rows, Some(&clock_for_query))) + .await + .map_err(|e| e.to_string())? + }, + timeout, + progress_clock, + None, + timeout_error, + ) + .await +} + +fn execute_query_blocking( + pool: &SqliteHandle, + sql: &str, + max_rows: Option, + progress_clock: Option<&crate::query::StreamProgressClock>, +) -> Result { let start = Instant::now(); let row_limit = query_result_row_limit(max_rows); @@ -2849,6 +2915,9 @@ fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option, ) -> Result { let row_limit = query_result_row_limit(max_rows); let mut remaining_offset = result_offset; @@ -724,6 +725,9 @@ async fn collect_first_result_limited( let mut truncated = false; while let Some(item) = stream.try_next().await.map_err(|e| e.to_string())? { + if let Some(progress_clock) = progress_clock { + progress_clock.mark(); + } match item { QueryItem::Metadata(metadata) if metadata.result_index() == 0 => { columns = columns_from_metadata(&metadata); @@ -3160,6 +3164,52 @@ pub async fn execute_query_with_max_rows( client: &mut SqlServerClient, sql: &str, max_rows: Option, +) -> Result { + execute_query_with_max_rows_inner(client, sql, max_rows, None).await +} + +/// Progress-aware variant of [`execute_query_with_max_rows`] for long transfers. +/// +/// Row-returning statements run under an inactivity budget: the clock is reset for +/// every streamed item, so a large table that keeps arriving is never cancelled +/// merely for exceeding the timeout in total — only a genuine stall is. Statements +/// without a result stream expose no incremental progress and keep the plain path. +pub(crate) async fn execute_query_with_max_rows_progress( + client: &mut SqlServerClient, + sql: &str, + max_rows: Option, + progress_clock: std::sync::Arc, + timeout: Option, +) -> Result { + let returns_rows = starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) + || sqlserver_dml_output_returns_rows(sql); + if !returns_rows { + // DDL/DML expose no incremental progress, so keep the original wall-clock + // budget: a hung write must still be bounded by the configured timeout. + return crate::query::wait_for_query_opt( + None, + timeout, + execute_query_with_max_rows_inner(client, sql, max_rows, None), + ) + .await; + } + let timeout_error = format!("Query timed out after {} seconds", timeout.map_or(0, |timeout| timeout.as_secs())); + let clock_for_query = progress_clock.clone(); + crate::query::await_stream_with_progress_timeout( + execute_query_with_max_rows_inner(client, sql, max_rows, Some(&clock_for_query)), + timeout, + progress_clock, + None, + timeout_error, + ) + .await +} + +async fn execute_query_with_max_rows_inner( + client: &mut SqlServerClient, + sql: &str, + max_rows: Option, + progress_clock: Option<&crate::query::StreamProgressClock>, ) -> Result { let start = Instant::now(); let result_offset = crate::query_result_sql::sqlserver_result_offset(sql); @@ -3175,8 +3225,16 @@ pub async fn execute_query_with_max_rows( }; let (result, messages) = capture_sqlserver_messages(async { let stream = sqlserver_driver_result(client.query(query.sql.as_str(), &[])).await?; - sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, sql, &query)) - .await + sqlserver_driver_result(collect_first_result_limited( + stream, + start, + max_rows, + result_offset, + sql, + &query, + progress_clock, + )) + .await }) .await; let mut result = query_result_with_server_messages(result?, messages); @@ -3264,6 +3322,7 @@ pub(crate) async fn execute_batch_with_max_rows_metadata( result_offset, sql, &query, + None, )) .await }) @@ -3358,7 +3417,8 @@ async fn execute_simple_batch_first_result_with_max_rows( let (result, messages) = capture_sqlserver_messages(async { let stream = sqlserver_driver_result(client.simple_query(sql)).await?; let query = SqlServerUnsafeTypeQuery::plain(sql); - sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, sql, &query)).await + sqlserver_driver_result(collect_first_result_limited(stream, start, max_rows, result_offset, sql, &query, None)) + .await }) .await; let mut result = query_result_with_server_messages(result?, messages); @@ -4327,9 +4387,8 @@ mod tests { let first_result = source.split("async fn execute_simple_batch_first_result_with_max_rows").nth(1).unwrap(); let first_result = first_result.split("fn strip_dbx_sqlserver_row_number_column").next().unwrap(); - assert!( - first_result.contains("collect_first_result_limited(stream, start, max_rows, result_offset, sql, &query)") - ); + assert!(first_result + .contains("collect_first_result_limited(stream, start, max_rows, result_offset, sql, &query, None)")); assert!(!first_result.contains("collect_result_sets_limited")); } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 6e7cde7871..ba085371e7 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -1470,15 +1470,33 @@ fn postgres_transaction_statement_error( pub(crate) struct StreamProgressClock { started_at: tokio::time::Instant, last_progress_ms: AtomicU64, + #[cfg(test)] + marked: std::sync::atomic::AtomicBool, } impl StreamProgressClock { pub(crate) fn new() -> Self { - Self { started_at: tokio::time::Instant::now(), last_progress_ms: AtomicU64::new(0) } + Self { + started_at: tokio::time::Instant::now(), + last_progress_ms: AtomicU64::new(0), + #[cfg(test)] + marked: std::sync::atomic::AtomicBool::new(false), + } } pub(crate) fn mark(&self) { self.last_progress_ms.store(self.started_at.elapsed().as_millis() as u64, Ordering::Relaxed); + #[cfg(test)] + self.marked.store(true, Ordering::Relaxed); + } + + /// Whether any progress has been recorded yet. Test-only: production code + /// only needs the derived inactivity window. A row read within the first + /// millisecond records a zero timestamp, so this cannot be derived from + /// `last_progress_ms`. + #[cfg(test)] + pub(crate) fn marked(&self) -> bool { + self.marked.load(Ordering::Relaxed) } fn elapsed_since_progress(&self) -> Duration { @@ -5988,6 +6006,48 @@ mod tests { assert_eq!(budget.query_timeout, None); } + #[tokio::test] + async fn stream_progress_timeout_survives_steady_progress_past_the_budget() { + // The timeout is an inactivity window, not a wall clock: a stream that keeps + // marking progress survives well past the budget, as long as each gap between + // marks is shorter than the timeout. + let clock = Arc::new(StreamProgressClock::new()); + let clock_for_rows = clock.clone(); + let result = await_stream_with_progress_timeout( + async move { + for _ in 0..20 { + tokio::time::sleep(Duration::from_millis(30)).await; + clock_for_rows.mark(); + } + Ok::<_, String>(42) + }, + Some(Duration::from_millis(200)), + clock, + None, + "timed out".to_string(), + ) + .await; + assert_eq!(result, Ok(42)); + } + + #[tokio::test] + async fn stream_progress_timeout_fires_when_no_progress_arrives() { + // A genuine stall — no progress for the whole budget — must still time out. + let clock = Arc::new(StreamProgressClock::new()); + let result = await_stream_with_progress_timeout( + async { + tokio::time::sleep(Duration::from_millis(600)).await; + Ok::<_, String>(42) + }, + Some(Duration::from_millis(200)), + clock, + None, + "timed out".to_string(), + ) + .await; + assert_eq!(result, Err("timed out".to_string())); + } + #[tokio::test] async fn sqlite_transaction_timeout_rolls_back_and_commits_nothing() { use std::sync::mpsc; diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 8d7e4ffcdd..ba58ffab8c 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -18,7 +18,7 @@ use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput}; use crate::query::{ agent_execute_query_params, is_dbx_query_timeout_error, pool_error_action, query_timeout_duration, - wait_for_query_opt, PoolErrorAction, QueryExecutionOptions, AGENT_PROTOCOL_MAX_ROWS, + wait_for_query_opt, PoolErrorAction, QueryExecutionOptions, StreamProgressClock, AGENT_PROTOCOL_MAX_ROWS, }; use crate::sql::{split_sql_statements, split_sql_statements_for_database}; use crate::sql_dialect::{ @@ -4854,7 +4854,19 @@ fn value_to_sql_literal(value: &serde_json::Value, _db_type: &DatabaseType) -> S /// dialects keep OFFSET paging (each page rescans and discards the rows before /// it, which is quadratic in table size) until their literal rules are audited. fn transfer_keyset_pagination_supported(db_type: &DatabaseType) -> bool { - matches!(db_type, DatabaseType::Postgres | DatabaseType::OpenGauss | DatabaseType::Gaussdb | DatabaseType::Kingbase) + matches!( + db_type, + DatabaseType::Postgres + | DatabaseType::OpenGauss + | DatabaseType::Gaussdb + | DatabaseType::Kingbase + | DatabaseType::Mysql + | DatabaseType::Doris + | DatabaseType::StarRocks + | DatabaseType::ManticoreSearch + | DatabaseType::Sqlite + | DatabaseType::SqlServer + ) } /// Column types whose keyset cursor value round-trips through a SQL text @@ -4894,6 +4906,77 @@ fn postgres_keyset_column_type_supported(data_type: &str) -> bool { SUPPORTED_PREFIXES.iter().any(|prefix| base.starts_with(prefix)) } +/// Dialect-aware keyset column-type gate. The Postgres family is covered by +/// [`postgres_keyset_column_type_supported`]; MySQL-family, SQLite and SQL Server +/// primary-key types that round-trip through `value_to_sql_literal` are enabled +/// here. Binary types stay OFF — their JSON form is lossy — so those keys keep +/// OFFSET paging. +fn keyset_column_type_supported(db_type: &DatabaseType, data_type: &str) -> bool { + match db_type { + DatabaseType::Mysql | DatabaseType::Doris | DatabaseType::StarRocks | DatabaseType::ManticoreSearch => { + mysql_keyset_column_type_supported(data_type) + } + DatabaseType::Sqlite => sqlite_keyset_column_type_supported(data_type), + DatabaseType::SqlServer => sqlserver_keyset_column_type_supported(data_type), + _ => postgres_keyset_column_type_supported(data_type), + } +} + +fn mysql_keyset_column_type_supported(data_type: &str) -> bool { + let base = data_type.trim().to_ascii_lowercase(); + let base = base.split('(').next().unwrap_or("").trim(); + const SUPPORTED: &[&str] = &[ + "int", + "integer", + "tinyint", + "smallint", + "mediumint", + "bigint", + "char", + "varchar", + "date", + "datetime", + "timestamp", + "year", + "decimal", + "numeric", + ]; + SUPPORTED.iter().any(|prefix| base.starts_with(prefix)) +} + +fn sqlite_keyset_column_type_supported(data_type: &str) -> bool { + let base = data_type.trim().to_ascii_lowercase(); + let base = base.split('(').next().unwrap_or("").trim(); + const SUPPORTED: &[&str] = &["int", "integer", "text", "char", "varchar", "character", "numeric", "decimal"]; + SUPPORTED.iter().any(|prefix| base.starts_with(prefix)) +} + +fn sqlserver_keyset_column_type_supported(data_type: &str) -> bool { + let base = data_type.trim().to_ascii_lowercase(); + let base = base.split('(').next().unwrap_or("").trim(); + const SUPPORTED: &[&str] = &[ + "int", + "bigint", + "smallint", + "tinyint", + "char", + "varchar", + "nchar", + "nvarchar", + "uniqueidentifier", + "date", + "datetime", + "datetime2", + "smalldatetime", + "time", + "decimal", + "numeric", + "money", + "smallmoney", + ]; + SUPPORTED.iter().any(|prefix| base.starts_with(prefix)) +} + /// Resolves the source primary key columns to their positions in the selected /// column list. Returns None — meaning the read loop keeps OFFSET paging — when /// the dialect is not keyset-capable, when a key column is not among the @@ -4911,7 +4994,7 @@ fn transfer_keyset_column_indexes( .iter() .map(|pk| { let index = columns.iter().position(|column| column.name == *pk)?; - postgres_keyset_column_type_supported(&columns[index].data_type).then_some(index) + keyset_column_type_supported(db_type, &columns[index].data_type).then_some(index) }) .collect() } @@ -5492,24 +5575,47 @@ async fn execute_on_pool_once( let pool_handle = state.pool_handle(pool_key).await; let pool = pool_handle.as_ref().ok_or("Connection not found")?; + // Transfer reads run under the per-connection operation budget. Drivers that + // expose an incremental result stream (MySQL, PostgreSQL, SQLite, SQL Server) + // use a *progress-aware* budget: the configured query timeout is an inactivity + // window reset for every row the server delivers, so transferring a large + // table is no longer cancelled just for exceeding the timeout in total. Drivers + // whose protocol returns the whole result in one shot — ClickHouse, InfluxDB, + // the Agent/JDBC path, external drivers and the DuckDB sidecar worker — expose + // no incremental progress, so they keep the plain wall-clock timeout. let result = match pool { PoolKind::Mysql(p, mode) => { let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; - wait_for_query_opt( - None, + // Row-returning reads run under a progress-aware budget: the timeout + // resets for every row the server delivers, so a large table is no + // longer cancelled just for taking longer than the timeout overall. + let progress_clock = Arc::new(StreamProgressClock::new()); + db::mysql::execute_query_with_max_rows_progress( + &p, + sql, + bare, + max_rows, + Default::default(), + progress_clock, query_timeout, - db::mysql::execute_query_with_max_rows(&p, sql, bare, max_rows, Default::default()), ) .await } PoolKind::Postgres(p) => { let p = p.clone(); - wait_for_query_opt(None, query_timeout, db::postgres::execute_query_with_max_rows(&p, sql, max_rows)).await + // Row-returning reads — the paging SELECTs a transfer issues — run + // under the driver's progress-aware budget: the configured query + // timeout becomes an inactivity window reset by every row the server + // delivers, so a large table is no longer cancelled just for taking + // longer than the timeout in total. + let progress_clock = Arc::new(StreamProgressClock::new()); + db::postgres::execute_query_with_max_rows_progress(&p, sql, max_rows, progress_clock, query_timeout).await } PoolKind::Sqlite(p) => { let p = p.clone(); - wait_for_query_opt(None, query_timeout, db::sqlite::execute_query_with_max_rows(&p, sql, max_rows)).await + let progress_clock = Arc::new(StreamProgressClock::new()); + db::sqlite::execute_query_with_max_rows_progress(&p, sql, max_rows, progress_clock, query_timeout).await } PoolKind::ClickHouse(client) => { let client = client.clone(); @@ -5539,10 +5645,16 @@ async fn execute_on_pool_once( PoolKind::SqlServer(client) => { let client = client.clone(); let mut client = client.lock().await; - let result = wait_for_query_opt( - None, + // Row-returning reads use the driver's progress-aware budget (see the + // SQL Server driver): a long but steady stream is never cancelled just + // for exceeding the timeout in total. + let progress_clock = Arc::new(StreamProgressClock::new()); + let result = db::sqlserver::execute_query_with_max_rows_progress( + &mut client, + sql, + max_rows, + progress_clock, query_timeout, - db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows), ) .await; drop(client); @@ -7596,7 +7708,11 @@ where request.source_catalog.as_deref(), ) }; - let result = execute_on_pool(state, source_pool_key, &sql).await?; + // Cap the result at `batch_size` (not the 10k default row limit): the + // paging SELECT is already `LIMIT batch_size`, and the loop below treats a + // short page as the last page. Capping lower than `batch_size` would make a + // large batch look short and truncate the transfer early. + let result = execute_on_pool_with_max_rows(state, source_pool_key, &sql, Some(batch_size)).await?; if let Some(indexes) = keyset_indexes.as_deref() { match advance_keyset_cursor(&mut keyset_cursor, &result.rows, indexes, table)? { KeysetAdvance::Advanced => {} @@ -8815,7 +8931,12 @@ where }; let (sql, mysql_spatial_markers) = mysql_spatial_transfer_select_sql(sql, &col_names, &col_types, source_db_type, target_db_type); - (execute_on_pool(state, source_pool_key, &sql).await?, mysql_spatial_markers) + // Cap the result at `batch_size` (not the 10k default row limit), so a + // large batch is never truncated into looking like a short final page. + ( + execute_on_pool_with_max_rows(state, source_pool_key, &sql, Some(batch_size)).await?, + mysql_spatial_markers, + ) }; let has_more = result.has_more; let row_count = result.rows.len(); @@ -13215,13 +13336,39 @@ mod tests { assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::OpenGauss), Some(vec![0])); assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::Gaussdb), Some(vec![0])); assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::Kingbase), Some(vec![0])); + assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::Mysql), Some(vec![0])); + assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::Sqlite), Some(vec![0])); + assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::SqlServer), Some(vec![0])); // Dialects whose cursor literal rendering is not audited keep OFFSET paging. - assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::Mysql), None); - assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::SqlServer), None); + assert_eq!(transfer_keyset_column_indexes(&columns, &pks, &DatabaseType::ClickHouse), None); // No primary key → no keyset cursor. assert_eq!(transfer_keyset_column_indexes(&columns, &[], &DatabaseType::Postgres), None); } + #[test] + fn mysql_sqlite_sqlserver_keyset_column_types_are_audited() { + // MySQL-family: integers, strings, dates and decimals round-trip; binary/blob stay OFF. + assert!(keyset_column_type_supported(&DatabaseType::Mysql, "int")); + assert!(keyset_column_type_supported(&DatabaseType::Mysql, "bigint unsigned")); + assert!(keyset_column_type_supported(&DatabaseType::Mysql, "varchar(64)")); + assert!(keyset_column_type_supported(&DatabaseType::Mysql, "datetime")); + assert!(!keyset_column_type_supported(&DatabaseType::Mysql, "binary(16)")); + assert!(!keyset_column_type_supported(&DatabaseType::Mysql, "varbinary(255)")); + assert!(!keyset_column_type_supported(&DatabaseType::Mysql, "blob")); + + // SQLite: integer/text; blob stays OFF. + assert!(keyset_column_type_supported(&DatabaseType::Sqlite, "INTEGER")); + assert!(keyset_column_type_supported(&DatabaseType::Sqlite, "TEXT")); + assert!(!keyset_column_type_supported(&DatabaseType::Sqlite, "BLOB")); + + // SQL Server: integers, strings, uniqueidentifier, dates; binary stays OFF. + assert!(keyset_column_type_supported(&DatabaseType::SqlServer, "int")); + assert!(keyset_column_type_supported(&DatabaseType::SqlServer, "nvarchar(64)")); + assert!(keyset_column_type_supported(&DatabaseType::SqlServer, "uniqueidentifier")); + assert!(keyset_column_type_supported(&DatabaseType::SqlServer, "datetime2")); + assert!(!keyset_column_type_supported(&DatabaseType::SqlServer, "varbinary(32)")); + } + #[test] fn keyset_column_indexes_require_selected_round_trippable_key_columns() { let columns = vec![ diff --git a/crates/dbx-core/src/transfer/rebuild_tests.rs b/crates/dbx-core/src/transfer/rebuild_tests.rs index e278941d2e..882e5e95e2 100644 --- a/crates/dbx-core/src/transfer/rebuild_tests.rs +++ b/crates/dbx-core/src/transfer/rebuild_tests.rs @@ -114,3 +114,60 @@ async fn transfer_rebuild_preview_plans_without_executing_ddl() { assert_eq!(rows.unwrap().rows[0][0], json!("original")); } } + +#[tokio::test] +async fn transfer_keyset_pagination_copies_every_row_across_many_batches() { + let directory = tempfile::tempdir().unwrap(); + let storage = crate::storage::Storage::open(&directory.path().join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new_with_plugin_dir(storage, directory.path().join("plugins"))); + for id in ["source", "target"] { + let path = directory.path().join(format!("{id}.db")); + crate::db::sqlite::connect_path_create_if_missing(path.to_str().unwrap()).await.unwrap(); + let config: ConnectionConfig = serde_json::from_value(json!({ + "id": id, "name": id, "db_type": "sqlite", "host": path.to_str().unwrap(), + "port": 0, "username": "", "password": "", "database": null, + "one_time": false, "save_password": false, "read_only": false + })) + .unwrap(); + state.configs.write().await.insert(id.to_string(), config); + } + let source_pool = ensure_transfer_pool(&state, "source", "main", None).await.unwrap(); + let target_pool = ensure_transfer_pool(&state, "target", "main", None).await.unwrap(); + + execute_on_pool(&state, &source_pool, "CREATE TABLE big(id INTEGER PRIMARY KEY, name TEXT)").await.unwrap(); + for id in 1..=25 { + execute_on_pool(&state, &source_pool, &format!("INSERT INTO big VALUES({id}, 'row-{id}')")).await.unwrap(); + } + + let request: TransferRequest = serde_json::from_value(json!({ + "transferId": uuid::Uuid::new_v4().to_string(), + "sourceConnectionId": "source", "sourceDatabase": "main", "sourceSchema": "main", + "targetConnectionId": "target", "targetDatabase": "main", "targetSchema": "main", + "tables": ["big"], "createTable": true, "content": "structureAndData", + "mode": "append", "batchSize": 3, "dropTargetBeforeCreate": false, + "dropTargetConfirmed": false + })) + .unwrap(); + + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &DatabaseType::Sqlite, + &DatabaseType::Sqlite, + &source_pool, + &target_pool, + &HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 25, "the whole table must be transferred"); + + let ids = execute_read_on_pool(&state, &target_pool, "SELECT id FROM big ORDER BY id").await.unwrap(); + let collected: Vec = ids.rows.iter().map(|row| row[0].as_i64().unwrap()).collect(); + assert_eq!(collected, (1..=25).collect::>(), "keyset pagination must not drop or duplicate rows"); +} diff --git a/crates/dbx-core/tests/live_mysql_transfer.rs b/crates/dbx-core/tests/live_mysql_transfer.rs index 051d914650..8bc8792c1c 100644 --- a/crates/dbx-core/tests/live_mysql_transfer.rs +++ b/crates/dbx-core/tests/live_mysql_transfer.rs @@ -1948,3 +1948,234 @@ async fn live_mysql_transfer_drop_target_retains_backup_on_failure() { setup_pool.disconnect().await.unwrap(); let _ = std::fs::remove_dir_all(dir); } + +#[tokio::test] +#[ignore = "requires disposable MySQL 8 endpoints via DBX_LIVE_MYSQL_TRANSFER_* variables"] +async fn live_mysql_keyset_pagination_copies_every_row() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_connection_id = format!("live-mysql-keyset-src-{suffix}"); + let target_connection_id = format!("live-mysql-keyset-dst-{suffix}"); + let source_database = format!("dbx_keyset_src_{}", &suffix[..12]); + let target_database = format!("dbx_keyset_dst_{}", &suffix[..12]); + let source_config = live_mysql_config(&source_connection_id); + let target_config = live_mysql_config(&target_connection_id); + let source_setup_pool = mysql::connect(&mysql_url(&source_config), Duration::from_secs(10)).await.unwrap(); + let target_setup_pool = mysql::connect(&mysql_url(&target_config), Duration::from_secs(10)).await.unwrap(); + + mysql::execute_query( + &source_setup_pool, + &format!( + "CREATE DATABASE `{source_database}` CHARACTER SET utf8mb4; \ + CREATE TABLE `{source_database}`.`big` (id BIGINT NOT NULL, name VARCHAR(64) NOT NULL, PRIMARY KEY (id)) ENGINE=InnoDB; \ + INSERT INTO `{source_database}`.`big` (id, name) \ + WITH RECURSIVE seq AS (SELECT 1 n UNION ALL SELECT n + 1 FROM seq WHERE n < 25) \ + SELECT n, CONCAT('row-', n) FROM seq" + ), + true, + ) + .await + .unwrap(); + mysql::execute_query( + &target_setup_pool, + &format!("CREATE DATABASE `{target_database}` CHARACTER SET utf8mb4"), + false, + ) + .await + .unwrap(); + + let task_tmp = std::path::PathBuf::from( + std::env::var("DBX_LIVE_MYSQL_TRANSFER_TMP_DIR") + .unwrap_or_else(|_| std::env::temp_dir().to_string_lossy().to_string()), + ); + let dir = task_tmp.join(format!("live-mysql-keyset-transfer-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + state.configs.write().await.insert(source_connection_id.clone(), source_config); + state.configs.write().await.insert(target_connection_id.clone(), target_config); + let source_pool_key = state.get_or_create_pool(&source_connection_id, Some(&source_database)).await.unwrap(); + let target_pool_key = state.get_or_create_pool(&target_connection_id, Some(&target_database)).await.unwrap(); + + let request = TransferRequest { + transfer_id: format!("live-mysql-keyset-{suffix}"), + source_connection_id, + source_database: source_database.clone(), + source_schema: source_database.clone(), + source_catalog: None, + target_connection_id, + target_database: target_database.clone(), + target_schema: target_database.clone(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 3, + }; + + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &DatabaseType::Mysql, + &DatabaseType::Mysql, + &source_pool_key, + &target_pool_key, + &std::collections::HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 25, "keyset pagination must copy every row"); + + let rows = mysql::execute_query( + &target_setup_pool, + &format!("SELECT id FROM `{target_database}`.`big` ORDER BY id"), + false, + ) + .await + .unwrap(); + let collected: Vec = rows + .rows + .iter() + .map(|row| { + let value = &row[0]; + value.as_i64().or_else(|| value.as_str().and_then(|s| s.parse().ok())).unwrap() + }) + .collect(); + assert_eq!(collected, (1..=25).collect::>(), "keyset pagination must not drop or duplicate rows"); + + let source_cleanup = + mysql::execute_query(&source_setup_pool, &format!("DROP DATABASE `{source_database}`"), false).await; + let target_cleanup = + mysql::execute_query(&target_setup_pool, &format!("DROP DATABASE `{target_database}`"), false).await; + source_setup_pool.disconnect().await.unwrap(); + target_setup_pool.disconnect().await.unwrap(); + let _ = std::fs::remove_dir_all(dir); + source_cleanup.unwrap(); + target_cleanup.unwrap(); +} + +#[tokio::test] +#[ignore = "requires disposable MySQL 8 endpoints via DBX_LIVE_MYSQL_TRANSFER_* variables"] +async fn live_mysql_progress_read_survives_total_duration_beyond_timeout() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_connection_id = format!("live-mysql-progress-src-{suffix}"); + let target_connection_id = format!("live-mysql-progress-dst-{suffix}"); + let source_database = format!("dbx_progress_src_{}", &suffix[..12]); + let target_database = format!("dbx_progress_dst_{}", &suffix[..12]); + let source_config = live_mysql_config(&source_connection_id); + let target_config = live_mysql_config(&target_connection_id); + let source_setup_pool = mysql::connect(&mysql_url(&source_config), Duration::from_secs(10)).await.unwrap(); + let target_setup_pool = mysql::connect(&mysql_url(&target_config), Duration::from_secs(10)).await.unwrap(); + + // The transfer spans several pages whose total duration far exceeds the 1s + // budget. With the timeout treated as an inactivity window, a steady stream must + // never be cancelled just for taking longer than the timeout in total. + mysql::execute_query( + &source_setup_pool, + &format!( + "SET SESSION cte_max_recursion_depth = 100000; \ + CREATE DATABASE `{source_database}` CHARACTER SET utf8mb4; \ + CREATE TABLE `{source_database}`.`big` (id BIGINT NOT NULL, name VARCHAR(200) NOT NULL, PRIMARY KEY (id)) ENGINE=InnoDB; \ + INSERT INTO `{source_database}`.`big` (id, name) \ + WITH RECURSIVE seq AS (SELECT 1 n UNION ALL SELECT n + 1 FROM seq WHERE n < 50000) \ + SELECT n, REPEAT('x', 200) FROM seq" + ), + true, + ) + .await + .unwrap(); + mysql::execute_query( + &target_setup_pool, + &format!("CREATE DATABASE `{target_database}` CHARACTER SET utf8mb4"), + false, + ) + .await + .unwrap(); + + let task_tmp = std::path::PathBuf::from( + std::env::var("DBX_LIVE_MYSQL_TRANSFER_TMP_DIR") + .unwrap_or_else(|_| std::env::temp_dir().to_string_lossy().to_string()), + ); + let dir = task_tmp.join(format!("live-mysql-progress-transfer-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + // The read runs under the source's query timeout; keep it at 1s so the test only + // passes when the transfer treats it as an inactivity budget, not a wall clock. + let mut source_config_short = source_config.clone(); + source_config_short.query_timeout_secs = 1; + state.configs.write().await.insert(source_connection_id.clone(), source_config_short); + state.configs.write().await.insert(target_connection_id.clone(), target_config); + let source_pool_key = state.get_or_create_pool(&source_connection_id, Some(&source_database)).await.unwrap(); + let target_pool_key = state.get_or_create_pool(&target_connection_id, Some(&target_database)).await.unwrap(); + + let request = TransferRequest { + transfer_id: format!("live-mysql-progress-{suffix}"), + source_connection_id, + source_database: source_database.clone(), + source_schema: source_database.clone(), + source_catalog: None, + target_connection_id, + target_database: target_database.clone(), + target_schema: target_database.clone(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 10000, + }; + + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &DatabaseType::Mysql, + &DatabaseType::Mysql, + &source_pool_key, + &target_pool_key, + &std::collections::HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 50000, "the whole table must be transferred across multiple pages without timing out"); + + let count = + mysql::execute_query(&target_setup_pool, &format!("SELECT COUNT(*) FROM `{target_database}`.`big`"), false) + .await + .unwrap(); + let count_value = &count.rows[0][0]; + let count_num = count_value.as_i64().or_else(|| count_value.as_str().and_then(|s| s.parse().ok())).unwrap(); + assert_eq!(count_num, 50000, "no row may be dropped or duplicated"); + + let source_cleanup = + mysql::execute_query(&source_setup_pool, &format!("DROP DATABASE `{source_database}`"), false).await; + let target_cleanup = + mysql::execute_query(&target_setup_pool, &format!("DROP DATABASE `{target_database}`"), false).await; + source_setup_pool.disconnect().await.unwrap(); + target_setup_pool.disconnect().await.unwrap(); + let _ = std::fs::remove_dir_all(dir); + source_cleanup.unwrap(); + target_cleanup.unwrap(); +} diff --git a/crates/dbx-core/tests/live_postgres_transfer.rs b/crates/dbx-core/tests/live_postgres_transfer.rs index e27e8294f1..e230bbc257 100644 --- a/crates/dbx-core/tests/live_postgres_transfer.rs +++ b/crates/dbx-core/tests/live_postgres_transfer.rs @@ -1783,3 +1783,324 @@ async fn live_postgres_transfer_rebuild_rejects_target_only_view_before_any_rena "preflight rejection must preserve both selected tables and the target-only view's result" ); } + +#[tokio::test] +#[ignore = "requires PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"] +async fn live_postgres_keyset_pagination_copies_every_row() { + let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL"); + let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone()); + let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap(); + let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap(); + let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_schema = format!("dbx_src_keyset_{}", &suffix[..8]); + let target_schema = format!("dbx_dst_keyset_{}", &suffix[..8]); + + postgres::execute_batch( + &source_pool, + &[ + format!("CREATE SCHEMA \"{source_schema}\""), + format!("CREATE TABLE \"{source_schema}\".\"big\" (\"id\" bigint PRIMARY KEY, \"name\" text NOT NULL)"), + format!("INSERT INTO \"{source_schema}\".\"big\" SELECT g, 'row-' || g FROM generate_series(1, 25) g"), + ], + ) + .await + .unwrap(); + postgres::execute_batch(&target_pool, &[format!("CREATE SCHEMA \"{target_schema}\"")]).await.unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-live-pg-keyset-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let source_connection_id = "live-pg-keyset-source"; + let target_connection_id = "live-pg-keyset-target"; + let source_pool_key = format!("{source_connection_id}:{source_database}"); + let target_pool_key = format!("{target_connection_id}:{target_database}"); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; + state + .configs + .write() + .await + .insert(source_connection_id.to_string(), postgres_test_config(source_connection_id, &source_database)); + state + .configs + .write() + .await + .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database)); + + let request = TransferRequest { + transfer_id: format!("live-pg-keyset-{suffix}"), + source_connection_id: source_connection_id.to_string(), + source_database: source_database.clone(), + source_schema: source_schema.clone(), + source_catalog: None, + target_connection_id: target_connection_id.to_string(), + target_database: target_database.clone(), + target_schema: target_schema.clone(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 3, + }; + let source_db_type = get_db_type(&state, source_connection_id).await.unwrap(); + let target_db_type = get_db_type(&state, target_connection_id).await.unwrap(); + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &source_db_type, + &target_db_type, + &source_pool_key, + &target_pool_key, + &std::collections::HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 25, "keyset pagination must copy every row"); + + let rows = postgres::execute_query( + &target_pool, + &format!("SELECT \"id\" FROM \"{target_schema}\".\"big\" ORDER BY \"id\""), + ) + .await + .unwrap(); + let collected: Vec = rows.rows.iter().map(|row| row[0].as_i64().unwrap()).collect(); + assert_eq!(collected, (1..=25).collect::>(), "keyset pagination must not drop or duplicate rows"); + + postgres::execute_batch(&source_pool, &[format!("DROP SCHEMA \"{source_schema}\" CASCADE")]).await.unwrap(); + postgres::execute_batch(&target_pool, &[format!("DROP SCHEMA \"{target_schema}\" CASCADE")]).await.unwrap(); + let _ = std::fs::remove_dir_all(dir); +} + +#[tokio::test] +#[ignore = "requires PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"] +async fn live_postgres_progress_read_survives_total_duration_beyond_timeout() { + let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL"); + let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone()); + let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap(); + let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap(); + let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_schema = format!("dbx_src_progress_{}", &suffix[..8]); + let target_schema = format!("dbx_dst_progress_{}", &suffix[..8]); + + // The transfer spans several pages whose total duration far exceeds the 1s + // budget. With the timeout treated as an inactivity window, a steady stream must + // never be cancelled just for taking longer than the timeout in total. + postgres::execute_batch( + &source_pool, + &[ + format!("CREATE SCHEMA \"{source_schema}\""), + format!("CREATE TABLE \"{source_schema}\".\"big\" (\"id\" bigint PRIMARY KEY, \"name\" text NOT NULL)"), + format!( + "INSERT INTO \"{source_schema}\".\"big\" SELECT g, repeat('x', 200) FROM generate_series(1, 50000) g" + ), + ], + ) + .await + .unwrap(); + postgres::execute_batch(&target_pool, &[format!("CREATE SCHEMA \"{target_schema}\"")]).await.unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-live-pg-progress-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let source_connection_id = "live-pg-progress-source"; + let target_connection_id = "live-pg-progress-target"; + let source_pool_key = format!("{source_connection_id}:{source_database}"); + let target_pool_key = format!("{target_connection_id}:{target_database}"); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; + // The read runs under the source's query timeout; keep it at 1s so the test only + // passes when the transfer treats it as an inactivity budget, not a wall clock. + let mut source_config = postgres_test_config(source_connection_id, &source_database); + source_config.query_timeout_secs = 1; + state.configs.write().await.insert(source_connection_id.to_string(), source_config); + state + .configs + .write() + .await + .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database)); + + let request = TransferRequest { + transfer_id: format!("live-pg-progress-{suffix}"), + source_connection_id: source_connection_id.to_string(), + source_database: source_database.clone(), + source_schema: source_schema.clone(), + source_catalog: None, + target_connection_id: target_connection_id.to_string(), + target_database: target_database.clone(), + target_schema: target_schema.clone(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 10000, + }; + let source_db_type = get_db_type(&state, source_connection_id).await.unwrap(); + let target_db_type = get_db_type(&state, target_connection_id).await.unwrap(); + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &source_db_type, + &target_db_type, + &source_pool_key, + &target_pool_key, + &std::collections::HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 50000, "the whole table must be transferred across multiple pages without timing out"); + + let count = postgres::execute_query(&target_pool, &format!("SELECT count(*) FROM \"{target_schema}\".\"big\"")) + .await + .unwrap(); + assert_eq!(count.rows[0][0].as_i64(), Some(50000), "no row may be dropped or duplicated"); + + postgres::execute_batch(&source_pool, &[format!("DROP SCHEMA \"{source_schema}\" CASCADE")]).await.unwrap(); + postgres::execute_batch(&target_pool, &[format!("DROP SCHEMA \"{target_schema}\" CASCADE")]).await.unwrap(); + let _ = std::fs::remove_dir_all(dir); +} + +#[tokio::test] +#[ignore = "requires PostgreSQL URLs via DBX_LIVE_PG_TRANSFER_SOURCE_URL and DBX_LIVE_PG_TRANSFER_TARGET_URL"] +async fn live_postgres_keyset_large_batch_copies_every_row() { + let source_url = std::env::var("DBX_LIVE_PG_TRANSFER_SOURCE_URL").expect("DBX_LIVE_PG_TRANSFER_SOURCE_URL"); + let target_url = std::env::var("DBX_LIVE_PG_TRANSFER_TARGET_URL").unwrap_or_else(|_| source_url.clone()); + let source_pool = postgres::connect(&source_url, std::time::Duration::from_secs(5)).await.unwrap(); + let target_pool = postgres::connect(&target_url, std::time::Duration::from_secs(5)).await.unwrap(); + let source_database = query_scalar(&source_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + let target_database = query_scalar(&target_pool, "SELECT current_database()").await.as_str().unwrap().to_string(); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_schema = format!("dbx_src_largebatch_{}", &suffix[..8]); + let target_schema = format!("dbx_dst_largebatch_{}", &suffix[..8]); + + // A batch larger than the 10k default row limit must not be truncated: the read + // must return every row, so the paging loop does not mistake a short page for the + // last page and stop early. + postgres::execute_batch( + &source_pool, + &[ + format!("CREATE SCHEMA \"{source_schema}\""), + format!("CREATE TABLE \"{source_schema}\".\"big\" (\"id\" bigint PRIMARY KEY, \"name\" text NOT NULL)"), + format!("INSERT INTO \"{source_schema}\".\"big\" SELECT g, 'row-' || g FROM generate_series(1, 50000) g"), + ], + ) + .await + .unwrap(); + postgres::execute_batch(&target_pool, &[format!("CREATE SCHEMA \"{target_schema}\"")]).await.unwrap(); + + let dir = std::env::temp_dir().join(format!("dbx-live-pg-largebatch-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = Arc::new(AppState::new(storage)); + let source_connection_id = "live-pg-largebatch-source"; + let target_connection_id = "live-pg-largebatch-target"; + let source_pool_key = format!("{source_connection_id}:{source_database}"); + let target_pool_key = format!("{target_connection_id}:{target_database}"); + state + .update_connection_pools(|connections| { + connections.insert(source_pool_key.clone(), PoolKind::Postgres(source_pool.clone())); + connections.insert(target_pool_key.clone(), PoolKind::Postgres(target_pool.clone())); + }) + .await; + state + .configs + .write() + .await + .insert(source_connection_id.to_string(), postgres_test_config(source_connection_id, &source_database)); + state + .configs + .write() + .await + .insert(target_connection_id.to_string(), postgres_test_config(target_connection_id, &target_database)); + + let request = TransferRequest { + transfer_id: format!("live-pg-largebatch-{suffix}"), + source_connection_id: source_connection_id.to_string(), + source_database: source_database.clone(), + source_schema: source_schema.clone(), + source_catalog: None, + target_connection_id: target_connection_id.to_string(), + target_database: target_database.clone(), + target_schema: target_schema.clone(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 50000, + }; + let source_db_type = get_db_type(&state, source_connection_id).await.unwrap(); + let target_db_type = get_db_type(&state, target_connection_id).await.unwrap(); + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &source_db_type, + &target_db_type, + &source_pool_key, + &target_pool_key, + &std::collections::HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await + .unwrap(); + assert_eq!(transferred, 50000, "a batch larger than the default row limit must not truncate the transfer"); + + let count = postgres::execute_query(&target_pool, &format!("SELECT count(*) FROM \"{target_schema}\".\"big\"")) + .await + .unwrap(); + assert_eq!(count.rows[0][0].as_i64(), Some(50000), "no row may be dropped"); + + postgres::execute_batch(&source_pool, &[format!("DROP SCHEMA \"{source_schema}\" CASCADE")]).await.unwrap(); + postgres::execute_batch(&target_pool, &[format!("DROP SCHEMA \"{target_schema}\" CASCADE")]).await.unwrap(); + let _ = std::fs::remove_dir_all(dir); +} diff --git a/crates/dbx-core/tests/live_sqlserver_transfer.rs b/crates/dbx-core/tests/live_sqlserver_transfer.rs index 4d66d83795..9c63da32ff 100644 --- a/crates/dbx-core/tests/live_sqlserver_transfer.rs +++ b/crates/dbx-core/tests/live_sqlserver_transfer.rs @@ -283,3 +283,355 @@ async fn live_sqlserver_transfer_rebuild_releases_constraint_and_index_names() { cleanup.expect("drop rebuild databases"); test_result.unwrap(); } + +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at SQL Server"] +async fn live_sqlserver_keyset_pagination_copies_every_row() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_db = format!("dbx_keyset_src_{}", &suffix[..12]); + let target_db = format!("dbx_keyset_dst_{}", &suffix[..12]); + let source_connection_id = format!("live-sqlserver-keyset-src-{suffix}"); + let target_connection_id = format!("live-sqlserver-keyset-dst-{suffix}"); + + let mut master = sqlserver_connect("master").await; + dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!("CREATE DATABASE [{source_db}]; CREATE DATABASE [{target_db}];"), + ) + .await + .expect("create keyset databases"); + + let mut source_client = sqlserver_connect(&source_db).await; + dbx_core::db::sqlserver::execute_batch( + &mut source_client, + "CREATE TABLE dbo.big (id INT NOT NULL CONSTRAINT PK_big PRIMARY KEY, name NVARCHAR(64) NOT NULL); \ + ;WITH seq AS (SELECT 1 n UNION ALL SELECT n + 1 FROM seq WHERE n < 25) \ + INSERT INTO dbo.big (id, name) \ + SELECT n, CONCAT('row-', n) FROM seq OPTION (MAXRECURSION 0);", + ) + .await + .expect("create source table"); + + let dir = std::env::temp_dir().join(format!("dbx-live-sqlserver-keyset-{suffix}")); + std::fs::create_dir_all(&dir).expect("create keyset directory"); + let storage = Storage::open(&dir.join("storage.db")).await.expect("open keyset storage"); + let state = Arc::new(AppState::new(storage)); + state + .configs + .write() + .await + .insert(source_connection_id.clone(), live_sqlserver_config(&source_connection_id, &source_db)); + state + .configs + .write() + .await + .insert(target_connection_id.clone(), live_sqlserver_config(&target_connection_id, &target_db)); + let source_pool_key = state.get_or_create_pool(&source_connection_id, Some(&source_db)).await.expect("source pool"); + let target_pool_key = state.get_or_create_pool(&target_connection_id, Some(&target_db)).await.expect("target pool"); + + let request = TransferRequest { + transfer_id: format!("live-sqlserver-keyset-{suffix}"), + source_connection_id: source_connection_id.clone(), + source_database: source_db.clone(), + source_schema: "dbo".to_string(), + source_catalog: None, + target_connection_id: target_connection_id.clone(), + target_database: target_db.clone(), + target_schema: "dbo".to_string(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 3, + }; + + let test_result = async { + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &DatabaseType::SqlServer, + &DatabaseType::SqlServer, + &source_pool_key, + &target_pool_key, + &HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await?; + assert_eq!(transferred, 25, "keyset pagination must copy every row"); + + let mut target_client = sqlserver_connect(&target_db).await; + let rows = dbx_core::db::sqlserver::execute_query(&mut target_client, "SELECT id FROM dbo.big ORDER BY id") + .await + .expect("read target ids"); + let collected: Vec = rows.rows.iter().map(|row| row[0].as_i64().unwrap()).collect(); + assert_eq!(collected, (1..=25).collect::>(), "keyset pagination must not drop or duplicate rows"); + Ok::<_, String>(()) + } + .await; + + let cleanup = dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!( + "ALTER DATABASE [{source_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{source_db}]; \ + ALTER DATABASE [{target_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{target_db}];" + ), + ) + .await; + let _ = std::fs::remove_dir_all(dir); + cleanup.expect("drop keyset databases"); + test_result.unwrap(); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at SQL Server"] +async fn live_sqlserver_progress_read_survives_total_duration_beyond_timeout() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_db = format!("dbx_progress_src_{}", &suffix[..12]); + let target_db = format!("dbx_progress_dst_{}", &suffix[..12]); + let source_connection_id = format!("live-sqlserver-progress-src-{suffix}"); + let target_connection_id = format!("live-sqlserver-progress-dst-{suffix}"); + + let mut master = sqlserver_connect("master").await; + dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!("CREATE DATABASE [{source_db}]; CREATE DATABASE [{target_db}];"), + ) + .await + .expect("create progress databases"); + + // The transfer spans several pages whose total duration far exceeds the 1s + // budget. With the timeout treated as an inactivity window, a steady stream must + // never be cancelled just for taking longer than the timeout in total. Generate + // the rows with a non-recursive cross join so the fixture itself stays fast. + let mut source_client = sqlserver_connect(&source_db).await; + dbx_core::db::sqlserver::execute_batch( + &mut source_client, + "CREATE TABLE dbo.big (id INT NOT NULL CONSTRAINT PK_big PRIMARY KEY, name NVARCHAR(64) NOT NULL); \ + INSERT INTO dbo.big (id, name) \ + SELECT t.n, REPLICATE('x', 64) \ + FROM ( \ + SELECT a.n + b.n * 10 + c.n * 100 + d.n * 1000 + e.n * 10000 + 1 AS n \ + FROM (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) a(n) \ + CROSS JOIN (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) b(n) \ + CROSS JOIN (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) c(n) \ + CROSS JOIN (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) d(n) \ + CROSS JOIN (VALUES (0),(1),(2),(3),(4),(5),(6),(7),(8),(9)) e(n) \ + ) t WHERE t.n <= 20000;", + ) + .await + .expect("create source table"); + + let dir = std::env::temp_dir().join(format!("dbx-live-sqlserver-progress-{suffix}")); + std::fs::create_dir_all(&dir).expect("create progress directory"); + let storage = Storage::open(&dir.join("storage.db")).await.expect("open progress storage"); + let state = Arc::new(AppState::new(storage)); + // The read runs under the source's query timeout; keep it at 1s so the test only + // passes when the transfer treats it as an inactivity budget, not a wall clock. + let mut source_config = live_sqlserver_config(&source_connection_id, &source_db); + source_config.query_timeout_secs = 1; + state.configs.write().await.insert(source_connection_id.clone(), source_config); + state + .configs + .write() + .await + .insert(target_connection_id.clone(), live_sqlserver_config(&target_connection_id, &target_db)); + let source_pool_key = state.get_or_create_pool(&source_connection_id, Some(&source_db)).await.expect("source pool"); + let target_pool_key = state.get_or_create_pool(&target_connection_id, Some(&target_db)).await.expect("target pool"); + + let request = TransferRequest { + transfer_id: format!("live-sqlserver-progress-{suffix}"), + source_connection_id: source_connection_id.clone(), + source_database: source_db.clone(), + source_schema: "dbo".to_string(), + source_catalog: None, + target_connection_id: target_connection_id.clone(), + target_database: target_db.clone(), + target_schema: "dbo".to_string(), + target_catalog: None, + tables: vec!["big".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 10000, + }; + + let test_result = async { + let transferred = transfer_table( + &state, + &request, + "big", + 0, + &DatabaseType::SqlServer, + &DatabaseType::SqlServer, + &source_pool_key, + &target_pool_key, + &HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await?; + assert_eq!(transferred, 20000, "the whole table must be transferred across multiple pages without timing out"); + + let mut target_client = sqlserver_connect(&target_db).await; + let count = dbx_core::db::sqlserver::execute_query(&mut target_client, "SELECT COUNT(*) FROM dbo.big") + .await + .expect("count target rows"); + assert_eq!(count.rows[0][0].as_i64(), Some(20000), "no row may be dropped or duplicated"); + Ok::<_, String>(()) + } + .await; + + let cleanup = dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!( + "ALTER DATABASE [{source_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{source_db}]; \ + ALTER DATABASE [{target_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{target_db}];" + ), + ) + .await; + let _ = std::fs::remove_dir_all(dir); + cleanup.expect("drop progress databases"); + test_result.unwrap(); +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_SQLSERVER_HOST/PORT/USER/PASSWORD pointing at SQL Server"] +async fn live_sqlserver_keyset_uniqueidentifier_datetime2_composite_key() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let source_db = format!("dbx_typed_src_{}", &suffix[..12]); + let target_db = format!("dbx_typed_dst_{}", &suffix[..12]); + let source_connection_id = format!("live-sqlserver-typed-src-{suffix}"); + let target_connection_id = format!("live-sqlserver-typed-dst-{suffix}"); + + let mut master = sqlserver_connect("master").await; + dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!("CREATE DATABASE [{source_db}]; CREATE DATABASE [{target_db}];"), + ) + .await + .expect("create typed databases"); + + let mut source_client = sqlserver_connect(&source_db).await; + dbx_core::db::sqlserver::execute_batch( + &mut source_client, + "CREATE TABLE dbo.typed_key ( \ + uid UNIQUEIDENTIFIER NOT NULL, \ + ts DATETIME2(3) NOT NULL, \ + name NVARCHAR(32) NOT NULL, \ + CONSTRAINT PK_typed_key PRIMARY KEY (uid, ts) \ + ); \ + INSERT INTO dbo.typed_key (uid, ts, name) VALUES \ + ('00000000-0000-0000-0000-000000000001', '2024-01-01 00:00:00.000', N'a'), \ + ('00000000-0000-0000-0000-000000000002', '2024-01-02 00:00:00.000', N'b'), \ + ('00000000-0000-0000-0000-000000000003', '2024-01-03 00:00:00.000', N'c'), \ + ('00000000-0000-0000-0000-000000000004', '2024-01-04 00:00:00.000', N'd'), \ + ('00000000-0000-0000-0000-000000000005', '2024-01-05 00:00:00.000', N'e'), \ + ('00000000-0000-0000-0000-000000000006', '2024-01-06 00:00:00.000', N'f'), \ + ('00000000-0000-0000-0000-000000000007', '2024-01-07 00:00:00.000', N'g'), \ + ('00000000-0000-0000-0000-000000000008', '2024-01-08 00:00:00.000', N'h');", + ) + .await + .expect("create source typed table"); + + let dir = std::env::temp_dir().join(format!("dbx-live-sqlserver-typed-{suffix}")); + std::fs::create_dir_all(&dir).expect("create typed directory"); + let storage = Storage::open(&dir.join("storage.db")).await.expect("open typed storage"); + let state = Arc::new(AppState::new(storage)); + state + .configs + .write() + .await + .insert(source_connection_id.clone(), live_sqlserver_config(&source_connection_id, &source_db)); + state + .configs + .write() + .await + .insert(target_connection_id.clone(), live_sqlserver_config(&target_connection_id, &target_db)); + let source_pool_key = state.get_or_create_pool(&source_connection_id, Some(&source_db)).await.expect("source pool"); + let target_pool_key = state.get_or_create_pool(&target_connection_id, Some(&target_db)).await.expect("target pool"); + + let request = TransferRequest { + transfer_id: format!("live-sqlserver-typed-{suffix}"), + source_connection_id: source_connection_id.clone(), + source_database: source_db.clone(), + source_schema: "dbo".to_string(), + source_catalog: None, + target_connection_id: target_connection_id.clone(), + target_database: target_db.clone(), + target_schema: "dbo".to_string(), + target_catalog: None, + tables: vec!["typed_key".to_string()], + create_table: true, + drop_target_before_create: false, + drop_target_confirmed: false, + content: TransferContent::default(), + objects: Vec::new(), + mode: TransferMode::Append, + target_table_name_case: TransferTableNameCase::Preserve, + quote_target_column_names: true, + ownership_policy: TransferOwnershipPolicy::Preserve, + batch_size: 3, + }; + + let test_result = async { + let transferred = transfer_table( + &state, + &request, + "typed_key", + 0, + &DatabaseType::SqlServer, + &DatabaseType::SqlServer, + &source_pool_key, + &target_pool_key, + &HashMap::new(), + &mut Vec::new(), + None, + |_| {}, + ) + .await?; + assert_eq!(transferred, 8, "keyset pagination must copy every typed row"); + + let mut target_client = sqlserver_connect(&target_db).await; + let rows = dbx_core::db::sqlserver::execute_query( + &mut target_client, + "SELECT LOWER(CAST(uid AS VARCHAR(36))) FROM dbo.typed_key ORDER BY uid", + ) + .await + .expect("read target uids"); + let collected: Vec = rows.rows.iter().map(|row| row[0].as_str().unwrap().to_string()).collect(); + let expected: Vec = (1..=8).map(|i| format!("00000000-0000-0000-0000-{i:012}")).collect(); + assert_eq!(collected, expected, "uniqueidentifier + datetime2 keyset cursor must round-trip every row"); + Ok::<_, String>(()) + } + .await; + + let cleanup = dbx_core::db::sqlserver::execute_batch( + &mut master, + &format!( + "ALTER DATABASE [{source_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{source_db}]; \ + ALTER DATABASE [{target_db}] SET SINGLE_USER WITH ROLLBACK IMMEDIATE; DROP DATABASE [{target_db}];" + ), + ) + .await; + let _ = std::fs::remove_dir_all(dir); + cleanup.expect("drop typed databases"); + test_result.unwrap(); +}