diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 53a36a0076..79c6ffd43f 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -3325,6 +3325,7 @@ impl AppState { return false; } + let mut checked_mysql_pool = None; let stale = { let connections = self.connections.read().await; let Some(pool) = connections.get(pool_key) else { @@ -3333,6 +3334,7 @@ impl AppState { match pool { PoolKind::Mysql(pool, _) => { let pool = pool.clone(); + checked_mysql_pool = Some(pool.clone()); drop(connections); match db::mysql::checkout_mysql_conn(&pool, HEALTH_CHECK_POOL_ACQUIRE_TIMEOUT).await { // Pool saturation means active work, not a dead connection. Removing this pool would @@ -3624,16 +3626,75 @@ impl AppState { return false; } + if let Some(checked_pool) = checked_mysql_pool { + return self.remove_stale_mysql_pool_if_current(pool_key, &checked_pool).await; + } + + let removed = self.connections.write().await.remove(pool_key); + + let Some(pool) = removed else { + return false; + }; + self.stop_keepalive_task(pool_key).await; self.pool_activity.write().await.remove(pool_key); self.postgres_cancel_contexts.write().await.remove(pool_key); - let removed = self.connections.write().await.remove(pool_key); - if let Some(pool) = removed { - self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await; - true - } else { - false - } + self.pool_routing_control().close_pool_with_timeout(pool_key.to_string(), pool).await; + true + } + + async fn remove_stale_mysql_pool_if_current(&self, pool_key: &str, checked_pool: &db::mysql::MySqlPool) -> bool { + self.remove_stale_mysql_pool_if_current_inner( + pool_key, + checked_pool, + #[cfg(test)] + None, + ) + .await + } + + async fn remove_stale_mysql_pool_if_current_inner( + &self, + pool_key: &str, + checked_pool: &db::mysql::MySqlPool, + #[cfg(test)] cleanup_barriers: Option<( + std::sync::Arc, + std::sync::Arc, + )>, + ) -> bool { + let routing = self.pool_routing_control(); + let removed = loop { + let mut connections = self.connections.write().await; + let is_current = matches!( + connections.get(pool_key), + Some(PoolKind::Mysql(current, _)) if checked_pool.is_same_pool(current) + ); + if !is_current { + log::debug!( + "MySQL connection pool '{pool_key}' was replaced while its health check was running; keeping the current route" + ); + return false; + } + let Ok(mut activity) = self.pool_activity.try_write() else { + drop(connections); + tokio::task::yield_now().await; + continue; + }; + let removed = remove_mysql_pool_if_current(&mut connections, pool_key, checked_pool) + .expect("checked MySQL pool must remain current while routing is locked"); + #[cfg(test)] + if let Some((route_removed, continue_cleanup)) = cleanup_barriers.as_ref() { + route_removed.wait().await; + continue_cleanup.wait().await; + } + routing.stop_keepalive(pool_key); + activity.remove(pool_key); + break removed; + }; + + self.postgres_cancel_contexts.write().await.remove(pool_key); + routing.close_pool_with_timeout(pool_key.to_string(), removed).await; + true } pub async fn reconnect_pool(&self, connection_id: &str, database: Option<&str>) -> Result { @@ -5304,6 +5365,18 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind { } } +fn remove_mysql_pool_if_current( + connections: &mut HashMap, + pool_key: &str, + expected: &db::mysql::MySqlPool, +) -> Option { + let is_current = matches!( + connections.get(pool_key), + Some(PoolKind::Mysql(current, _)) if expected.is_same_pool(current) + ); + is_current.then(|| connections.remove(pool_key)).flatten() +} + async fn close_pool_kind(pool: PoolKind) -> Result<(), String> { match pool { PoolKind::Mysql(p, _) => { @@ -7334,6 +7407,102 @@ mod tests { assert_eq!(super::mysql_pool_max_connections_for_session(Some("tab-1")), 1); } + #[test] + fn stale_mysql_pool_observation_does_not_remove_replacement_generation() { + let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + let checked_clone = checked.clone(); + let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + assert!(checked.is_same_pool(&checked_clone)); + assert!(!checked.is_same_pool(&replacement)); + + let mut connections = std::collections::HashMap::from([( + "conn:app".to_string(), + PoolKind::Mysql(replacement.clone(), MysqlMode::Normal), + )]); + + assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &checked).is_none()); + assert!( + matches!(connections.get("conn:app"), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current)) + ); + + assert!(super::remove_mysql_pool_if_current(&mut connections, "conn:app", &replacement).is_some()); + assert!(!connections.contains_key("conn:app")); + } + + #[tokio::test] + async fn stale_mysql_cleanup_preserves_replacement_support_state() { + let (state, dir) = test_app_state().await; + let state = std::sync::Arc::new(state); + let pool_key = "conn:app"; + let checked = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + let replacement = crate::db::mysql::MySqlPool::new("mysql://root@127.0.0.1:3306/app", 10); + let mut config = mysql_config(Some("app")); + config.keepalive_interval_secs = 60; + state + .connections + .write() + .await + .insert(pool_key.to_string(), PoolKind::Mysql(checked.clone(), MysqlMode::Normal)); + state.pool_activity.write().await.insert(pool_key.to_string(), super::PoolActivity::now()); + state.start_keepalive_task( + pool_key, + &PoolKind::Mysql(checked.clone(), MysqlMode::Normal), + &config, + #[cfg(feature = "mq-admin")] + None, + ); + assert_eq!(state.supervised_task_count(), 1); + + let route_removed = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let continue_cleanup = std::sync::Arc::new(tokio::sync::Barrier::new(2)); + let cleanup_state = state.clone(); + let cleanup_checked = checked.clone(); + let cleanup_route_removed = route_removed.clone(); + let cleanup_continue = continue_cleanup.clone(); + let cleanup = tokio::spawn(async move { + cleanup_state + .remove_stale_mysql_pool_if_current_inner( + pool_key, + &cleanup_checked, + Some((cleanup_route_removed, cleanup_continue)), + ) + .await + }); + route_removed.wait().await; + + let publish_state = state.clone(); + let publish_config = config.clone(); + let publish_replacement = replacement.clone(); + let mut publish = tokio::spawn(async move { + publish_state + .insert_connection_pool( + pool_key.to_string(), + PoolKind::Mysql(publish_replacement, MysqlMode::Normal), + &publish_config, + ) + .await + }); + + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut publish).await.is_err(), + "replacement publication must wait while stale support state is cleaned" + ); + continue_cleanup.wait().await; + assert!(cleanup.await.unwrap()); + publish.await.unwrap().unwrap(); + + let connections = state.connections.read().await; + assert!( + matches!(connections.get(pool_key), Some(PoolKind::Mysql(current, _)) if replacement.is_same_pool(current)) + ); + drop(connections); + assert!(state.pool_activity.read().await.contains_key(pool_key)); + assert_eq!(state.supervised_task_count(), 1); + + state.shutdown(Duration::from_secs(1)).await; + let _ = std::fs::remove_dir_all(dir); + } + #[test] fn metadata_concurrency_reserves_query_capacity() { assert_eq!(super::metadata_concurrency_limit(DatabaseType::Mysql, 10), 6); diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 945c9cd888..25844956b9 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -3,7 +3,8 @@ use serde_json::Value; use std::collections::{HashMap, HashSet}; use std::future::Future; use std::io::{BufWriter, Write}; -use std::sync::RwLock; +use std::pin::Pin; +use std::sync::{Arc, RwLock}; use std::time::Duration; use tokio_util::sync::CancellationToken; @@ -1586,10 +1587,10 @@ pub async fn clear_export_cancelled(export_id: &str) { /// not accept a cancellation token; polling here keeps the export task /// responsive and dropping the pending future follows the same bounded /// prefetch cancellation behavior used below. -async fn await_export_operation(export_id: &str, operation: F) -> Result -where - F: Future>, -{ +async fn await_export_operation( + export_id: &str, + operation: Pin> + Send + '_>>, +) -> Result { tokio::pin!(operation); loop { if is_export_cancelled_now(export_id) { @@ -1609,6 +1610,43 @@ where } } +struct AbortExportTaskOnDrop(tokio::task::AbortHandle); + +impl Drop for AbortExportTaskOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +async fn get_export_table_ddl_isolated( + state: Arc, + connection_id: String, + database: String, + schema: String, + table: String, +) -> Result { + let task = tokio::spawn(async move { + crate::schema::get_table_relation_export_ddl_core(&state, &connection_id, &database, &schema, &table, None) + .await + }); + let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle()); + task.await.map_err(|error| format!("Database export metadata task failed: {error}"))? +} + +async fn get_export_table_columns_isolated( + state: Arc, + connection_id: String, + database: String, + schema: String, + table: String, +) -> Result, String> { + let task = tokio::spawn(async move { + crate::schema::get_columns_core(&state, &connection_id, &database, &schema, &table).await + }); + let _abort_on_drop = AbortExportTaskOnDrop(task.abort_handle()); + task.await.map_err(|error| format!("Database export metadata task failed: {error}"))? +} + fn snapshot_batch_cancelled(db_type: &DatabaseType, export_id: &str) -> bool { *db_type != DatabaseType::Postgres && is_export_cancelled_now(export_id) } @@ -2106,7 +2144,7 @@ mod windows_export_destination { } pub async fn export_database_sql_core( - state: &crate::connection::AppState, + state: &Arc, request: &DatabaseExportRequest, on_progress: impl Fn(ExportProgress) + Sync, ) -> Result<(), String> { @@ -2126,7 +2164,7 @@ pub async fn export_database_sql_core( } async fn export_database_sql_core_inner( - state: &crate::connection::AppState, + state: &Arc, request: &DatabaseExportRequest, on_progress: impl Fn(ExportProgress) + Sync, ) -> Result<(), String> { @@ -2275,7 +2313,7 @@ async fn export_database_sql_core_inner( let postgres_sequences = if request.include_structure && matches!(db_type, DatabaseType::Postgres) { match await_export_operation( &request.export_id, - list_postgres_export_sequences( + Box::pin(list_postgres_export_sequences( state, &pool_key, &request.schema, @@ -2283,7 +2321,7 @@ async fn export_database_sql_core_inner( &request.excluded_tables, request.include_objects, request.fail_on_error, - ), + )), ) .await { @@ -2416,8 +2454,8 @@ async fn export_database_sql_core_inner( use futures::StreamExt; let prefetch_targets: Vec<(usize, String)> = tables.iter().enumerate().map(|(index, table_info)| (index, table_info.name.clone())).collect(); - let mut prefetch_stream = - futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| async move { + let mut prefetch_stream = futures::stream::iter(prefetch_targets.into_iter().map(|(index, table_name)| { + Box::pin(async move { if is_export_cancelled_now(&request.export_id) { return (index, PrefetchedTableMetadata { ddl: None, columns: None }); } @@ -2425,14 +2463,13 @@ async fn export_database_sql_core_inner( Some( await_export_operation( &request.export_id, - crate::schema::get_table_relation_export_ddl_core( - state, - &request.connection_id, - &request.database, - &request.schema, - &table_name, - None, - ), + Box::pin(get_export_table_ddl_isolated( + state.clone(), + request.connection_id.clone(), + request.database.clone(), + request.schema.clone(), + table_name.clone(), + )), ) .await, ) @@ -2446,13 +2483,13 @@ async fn export_database_sql_core_inner( Some( await_export_operation( &request.export_id, - crate::schema::get_columns_core( - state, - &request.connection_id, - &request.database, - &request.schema, - &table_name, - ), + Box::pin(get_export_table_columns_isolated( + state.clone(), + request.connection_id.clone(), + request.database.clone(), + request.schema.clone(), + table_name.clone(), + )), ) .await, ) @@ -2460,8 +2497,9 @@ async fn export_database_sql_core_inner( None }; (index, PrefetchedTableMetadata { ddl, columns }) - })) - .buffer_unordered(database_export_metadata_prefetch_concurrency(db_type)); + }) + })) + .buffer_unordered(database_export_metadata_prefetch_concurrency(db_type)); while let Some((index, metadata)) = prefetch_stream.next().await { if metadata .ddl @@ -2597,14 +2635,13 @@ async fn export_database_sql_core_inner( None => { await_export_operation( &request.export_id, - crate::schema::get_table_relation_export_ddl_core( - state, - &request.connection_id, - &request.database, - &request.schema, - table_name, - None, - ), + Box::pin(get_export_table_ddl_isolated( + state.clone(), + request.connection_id.clone(), + request.database.clone(), + request.schema.clone(), + table_name.clone(), + )), ) .await } @@ -2645,13 +2682,13 @@ async fn export_database_sql_core_inner( None => { await_export_operation( &request.export_id, - crate::schema::get_columns_core( - state, - &request.connection_id, - &request.database, - &request.schema, - table_name, - ), + Box::pin(get_export_table_columns_isolated( + state.clone(), + request.connection_id.clone(), + request.database.clone(), + request.schema.clone(), + table_name.clone(), + )), ) .await } @@ -3124,10 +3161,13 @@ mod tests { clear_export_cancelled(&export_id).await; let task_export_id = export_id.clone(); let task = tokio::spawn(async move { - await_export_operation(&task_export_id, async { - tokio::time::sleep(std::time::Duration::from_secs(30)).await; - Ok::<_, String>(()) - }) + await_export_operation( + &task_export_id, + Box::pin(async { + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + Ok::<_, String>(()) + }), + ) .await }); @@ -3145,7 +3185,7 @@ mod tests { async fn await_export_operation_preserves_completed_metadata() { let export_id = format!("complete-metadata-{}", uuid::Uuid::new_v4()); clear_export_cancelled(&export_id).await; - let result = await_export_operation(&export_id, async { Ok::<_, String>(42_u32) }).await; + let result = await_export_operation(&export_id, Box::pin(async { Ok::<_, String>(42_u32) })).await; assert_eq!(result, Ok(42)); } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 3451e59b85..9bf583178d 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -55,6 +55,14 @@ impl MySqlPool { pub async fn disconnect(self) -> Result<(), mysql_async::Error> { self.inner.disconnect().await } + + /// Returns whether both handles refer to the same underlying pool + /// generation. Pool options are not an identity: a reconnect creates a new + /// pool with identical options, and a late health probe for the old pool + /// must not be allowed to remove that replacement from routing. + pub(crate) fn is_same_pool(&self, other: &Self) -> bool { + std::sync::Arc::ptr_eq(&self.inner.metrics(), &other.inner.metrics()) + } } impl Deref for MySqlPool { diff --git a/crates/dbx-core/tests/database_export_partition_ddl.rs b/crates/dbx-core/tests/database_export_partition_ddl.rs index 6cf0e758aa..fad4f28977 100644 --- a/crates/dbx-core/tests/database_export_partition_ddl.rs +++ b/crates/dbx-core/tests/database_export_partition_ddl.rs @@ -13,6 +13,7 @@ mod support; use std::process::Command; +use std::sync::Arc; use dbx_core::connection::AppState; use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest}; @@ -76,7 +77,7 @@ async fn run_database_export_of_partition_tree_has_no_duplicates_and_replays() { let dir = std::env::temp_dir().join(format!("dbx-export-partition-ddl-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); let connection_id = "export-partition-ddl-conn"; state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); @@ -234,7 +235,7 @@ async fn partition_tree_ddl_query_count_does_not_scale_with_partition_count() { let dir = std::env::temp_dir().join(format!("dbx-partition-query-count-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); let connection_id = "partition-query-count-conn"; state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); @@ -302,7 +303,7 @@ async fn display_ddl_error_distinguishes_wrong_relkind_from_missing_relation() { let dir = std::env::temp_dir().join(format!("dbx-partition-relkind-error-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); let connection_id = "partition-relkind-error-conn"; state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); diff --git a/crates/dbx-core/tests/database_export_prefetch.rs b/crates/dbx-core/tests/database_export_prefetch.rs index fa5480c312..16f63e5011 100644 --- a/crates/dbx-core/tests/database_export_prefetch.rs +++ b/crates/dbx-core/tests/database_export_prefetch.rs @@ -8,6 +8,7 @@ use std::process::Command; use dbx_core::connection::AppState; use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest}; use dbx_core::storage::Storage; +use std::sync::Arc; use support::{postgres_test_config, psql, start_docker_postgres}; #[test] @@ -54,7 +55,7 @@ async fn run_database_export_writes_structure_and_data_for_all_tables() { let dir = std::env::temp_dir().join(format!("dbx-export-prefetch-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); let connection_id = "export-prefetch-conn"; state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); diff --git a/crates/dbx-core/tests/live_manual_transaction.rs b/crates/dbx-core/tests/live_manual_transaction.rs index c2c88e0d50..5cbf5c368b 100644 --- a/crates/dbx-core/tests/live_manual_transaction.rs +++ b/crates/dbx-core/tests/live_manual_transaction.rs @@ -44,10 +44,10 @@ fn live_config(prefix: &str, db_type: DatabaseType, default_port: u16) -> Connec .expect("live connection config should deserialize") } -async fn app_state_with_config(config: ConnectionConfig) -> (AppState, std::path::PathBuf) { +async fn app_state_with_config(config: ConnectionConfig) -> (Arc, std::path::PathBuf) { let db_path = std::env::temp_dir().join(format!("dbx-live-manual-txn-{}.db", uuid::Uuid::new_v4().simple())); let storage = Storage::open(&db_path).await.expect("open temp storage"); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(config.id.clone(), config); (state, db_path) } diff --git a/crates/dbx-core/tests/live_mysql_database_export.rs b/crates/dbx-core/tests/live_mysql_database_export.rs index 4084f95ed6..d923e5b5db 100644 --- a/crates/dbx-core/tests/live_mysql_database_export.rs +++ b/crates/dbx-core/tests/live_mysql_database_export.rs @@ -5,6 +5,7 @@ use dbx_core::query::execute_sql_statement; use dbx_core::sql::SqlFileRequest; use dbx_core::sql_file_import::execute_sql_file_path; use dbx_core::storage::Storage; +use std::sync::Arc; use tokio_util::sync::CancellationToken; fn exported_view_position(sql: &str, view_name: &str) -> Option { @@ -47,7 +48,7 @@ async fn live_mysql_database_export_restores_dependent_views() { let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-{suffix}")); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(connection_id.clone(), live_mysql_config(&connection_id)); for sql in [ @@ -134,6 +135,119 @@ async fn live_mysql_database_export_restores_dependent_views() { assert_eq!(spatial_result.rows, vec![vec![serde_json::json!(4326), serde_json::json!(1)]]); } +/// Regression coverage for #6882. A whole-database export prefetches MySQL +/// metadata concurrently, so late health checks for an older pool generation +/// must not remove a replacement pool from routing. Empty tables must still +/// contribute their DDL even though they naturally produce no INSERT rows. +#[test] +#[ignore = "requires a disposable MySQL endpoint"] +fn live_mysql_database_export_handles_many_tables_including_empty_tables() { + let handle = std::thread::Builder::new() + .name("live-mysql-export-many-tables".to_string()) + // Tokio worker threads use a 2 MiB stack by default. Keep this test at + // that production-sized boundary so large metadata futures cannot + // silently rely on the oversized stacks used by older export tests. + .stack_size(2 * 1024 * 1024) + .spawn(|| { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("build live MySQL export runtime") + .block_on(run_live_mysql_database_export_handles_many_tables_including_empty_tables()); + }) + .expect("spawn live MySQL export thread"); + if let Err(panic) = handle.join() { + std::panic::resume_unwind(panic); + } +} + +async fn run_live_mysql_database_export_handles_many_tables_including_empty_tables() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let connection_id = format!("live-mysql-export-many-tables-{suffix}"); + let database = format!("dbx_export_many_{suffix}"); + let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-many-tables-{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(connection_id.clone(), live_mysql_config(&connection_id)); + + execute_sql_statement(&state, &connection_id, "", &format!("CREATE DATABASE `{database}`"), None, None) + .await + .unwrap(); + for index in 1..=80 { + let table = format!("export_probe_{index:03}"); + let create = format!( + "CREATE TABLE `{database}`.`{table}` (\ + id INT PRIMARY KEY, parent_id INT NULL, name VARCHAR(120) NOT NULL, note TEXT, \ + created_at DATETIME, amount DECIMAL(12, 2), active BOOLEAN, metadata JSON, \ + INDEX idx_parent (parent_id))" + ); + execute_sql_statement(&state, &connection_id, "", &create, None, None).await.unwrap(); + if index % 5 != 0 { + let insert = format!( + "INSERT INTO `{database}`.`{table}` \ + (id, parent_id, name, note, created_at, amount, active, metadata) \ + VALUES ({index}, NULL, 'row-{index:03}', 'batch export probe', \ + '2026-08-25 12:00:00', 12.34, TRUE, '{{\"index\": {index}}}')" + ); + execute_sql_statement(&state, &connection_id, "", &insert, None, None).await.unwrap(); + } + } + let test_result = async { + for attempt in 1..=5 { + let file_path = dir.join(format!("export-{attempt}.sql")); + let request = DatabaseExportRequest { + export_id: format!("live-mysql-export-many-tables-{suffix}-{attempt}"), + connection_id: connection_id.clone(), + database: database.clone(), + schema: database.clone(), + file_path: file_path.to_string_lossy().to_string(), + selected_tables: Vec::new(), + 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: false, + snapshot_session_id: None, + batch_size: 1000, + }; + + export_database_sql_core(&state, &request, |_| {}).await?; + let exported = std::fs::read_to_string(&file_path).map_err(|error| error.to_string())?; + if exported.contains("-- ERROR") || exported.contains("Agent runtime is unavailable") { + return Err(format!("attempt {attempt} contained an inline export error:\n{exported}")); + } + for index in 1..=80 { + let table = format!("export_probe_{index:03}"); + let create = format!("CREATE TABLE `{table}`"); + if !exported.contains(&create) { + return Err(format!("attempt {attempt} did not export DDL for {table}")); + } + if index % 5 == 0 { + let qualified_insert = format!("INSERT INTO `{database}`.`{table}`"); + let unqualified_insert = format!("INSERT INTO `{table}`"); + if exported.contains(&qualified_insert) || exported.contains(&unqualified_insert) { + return Err(format!("attempt {attempt} unexpectedly exported rows for empty table {table}")); + } + } else if !exported.contains(&format!("'row-{index:03}'")) { + return Err(format!("attempt {attempt} did not export the seeded row for {table}")); + } + } + } + Ok::<_, String>(()) + } + .await; + + let cleanup = + execute_sql_statement(&state, &connection_id, "", &format!("DROP DATABASE `{database}`"), None, None).await; + cleanup.unwrap(); + std::fs::remove_dir_all(dir).unwrap(); + test_result.unwrap(); +} + /// Regression test for #6109 ("backup always errors"): the very first export /// to a destination directory that has never been configured or used before /// (a normal, not-yet-created local folder -- e.g. the user just typed a new @@ -157,7 +271,7 @@ async fn live_mysql_database_export_creates_missing_destination_directory() { let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-missing-dir-{suffix}")); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(connection_id.clone(), live_mysql_config(&connection_id)); for sql in [ @@ -222,7 +336,7 @@ async fn live_mysql_database_export_refuses_to_recreate_a_destination_that_disap let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-vanished-dir-{suffix}")); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(connection_id.clone(), live_mysql_config(&connection_id)); for sql in [ @@ -301,7 +415,7 @@ async fn live_mysql_database_export_refuses_a_destination_that_vanished_before_i let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-preconfigured-vanished-dir-{suffix}")); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(connection_id.clone(), live_mysql_config(&connection_id)); for sql in [ diff --git a/crates/dbx-core/tests/live_mysql_export_table_order.rs b/crates/dbx-core/tests/live_mysql_export_table_order.rs index 830b8f2fb0..83d4d66c48 100644 --- a/crates/dbx-core/tests/live_mysql_export_table_order.rs +++ b/crates/dbx-core/tests/live_mysql_export_table_order.rs @@ -3,6 +3,7 @@ use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest} use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; use dbx_core::query::execute_sql_statement; use dbx_core::storage::Storage; +use std::sync::Arc; fn live_mysql_config(id: &str) -> ConnectionConfig { let host = std::env::var("DBX_LIVE_SQL_FILE_MYSQL_HOST").expect("DBX_LIVE_SQL_FILE_MYSQL_HOST"); @@ -54,7 +55,7 @@ async fn live_mysql_database_export_table_order_is_not_alphabetical_when_fk_reor let dir = std::env::temp_dir().join(format!("dbx-live-mysql-export-order-{suffix}")); std::fs::create_dir_all(&dir).unwrap(); let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); - let state = AppState::new(storage); + let state = Arc::new(AppState::new(storage)); state.configs.write().await.insert(connection_id.clone(), live_mysql_config(&connection_id)); for sql in [