From 81f787e8d1795866c965d467b37315fe1cd33818 Mon Sep 17 00:00:00 2001 From: vrustx <279631638@qq.com> Date: Sun, 19 Jul 2026 08:50:27 +0800 Subject: [PATCH] feat(export): prefetch table metadata during database export --- crates/dbx-core/src/database_export.rs | 171 ++++++++++++-- .../tests/database_export_prefetch.rs | 210 ++++++++++++++++++ 2 files changed, 363 insertions(+), 18 deletions(-) create mode 100644 crates/dbx-core/tests/database_export_prefetch.rs diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 3a6d265b06..cf90088996 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -1057,6 +1057,18 @@ pub async fn begin_database_backup_snapshot_core( Ok(DatabaseBackupSnapshot { session_id, schemas }) } +/// 只有确认底层支持并发请求的多连接池才允许并发预取导出元数据。 +/// SqlServer（Arc 串行客户端）、Agent/ExternalDriver（插件请求超时覆盖 +/// 排队时间且超时会终止 sidecar）、SQLite/DuckDB 等单连接类型都必须回退串行。 +fn concurrent_metadata_prefetch_allowed(pool_kind: Option<&crate::connection::PoolKind>) -> bool { + matches!( + pool_kind, + Some(crate::connection::PoolKind::Postgres(_)) + | Some(crate::connection::PoolKind::Mysql(..)) + | Some(crate::connection::PoolKind::ClickHouse(_)) + ) +} + fn record_export_error(file: &mut std::fs::File, fail_on_error: bool, message: String) -> Result<(), String> { if fail_on_error { Err(message) @@ -1254,7 +1266,89 @@ pub async fn export_database_sql_core( object_index += 1; } - for table_info in &tables { + // 预取各表的 DDL 与列元数据：逐表串行往返在多表数据库上是整库导出耗时的 + // 主要来源（每表 1-2 次网络往返 × 表数）。有界并发预取后，下方写出循环仍按 + // 原顺序消费，文件内容与逐表查询完全一致。 + struct PrefetchedTableMetadata { + ddl: Option>, + columns: Option, String>>, + } + const EXPORT_METADATA_PREFETCH_CONCURRENCY: usize = 8; + let mut prefetched_table_metadata: Vec> = Vec::new(); + prefetched_table_metadata.resize_with(tables.len(), || None); + // 防护门按「实际连接池种类」放行，而非数据库类型的能力标记：只有确认底层 + // 支持并发请求的多连接池（Postgres/MySQL/ClickHouse）才启用并发预取。反例： + // 原生 SQL Server 是 Arc> 串行客户端；Agent/外部 JDBC（含 SQL Server + // legacy profile、PrestoSQL 等路由结果）的插件请求超时覆盖排队时间且超时会 + // 终止 sidecar；SQLite/DuckDB 等为单连接。被挡住的场景预取 Vec 保持全 None， + // 写出循环内的 None 回退路径即原有的逐表串行直查行为。 + let concurrent_prefetch_is_safe = + match state.get_or_create_pool(&request.connection_id, Some(&request.database)).await { + Ok(metadata_pool_key) => { + concurrent_metadata_prefetch_allowed(state.connections.read().await.get(&metadata_pool_key)) + } + // 建池失败时不预取，让写出循环的直查路径按原有方式报告错误 + Err(_) => false, + }; + if concurrent_prefetch_is_safe + && (request.include_structure || request.include_data) + && !tables.is_empty() + && !is_export_cancelled(&request.export_id).await + { + 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 { + if is_export_cancelled_now(&request.export_id) { + return (index, PrefetchedTableMetadata { ddl: None, columns: None }); + } + let ddl = if request.include_structure { + Some( + crate::schema::get_table_ddl_core( + state, + &request.connection_id, + &request.database, + &request.schema, + &table_name, + None, + ) + .await, + ) + } else { + None + }; + if is_export_cancelled_now(&request.export_id) { + return (index, PrefetchedTableMetadata { ddl, columns: None }); + } + let columns = if request.include_data { + Some( + crate::schema::get_columns_core( + state, + &request.connection_id, + &request.database, + &request.schema, + &table_name, + ) + .await, + ) + } else { + None + }; + (index, PrefetchedTableMetadata { ddl, columns }) + })) + .buffer_unordered(EXPORT_METADATA_PREFETCH_CONCURRENCY); + while let Some((index, metadata)) = prefetch_stream.next().await { + prefetched_table_metadata[index] = Some(metadata); + // 取消后不再调度新的预取任务（已在途的任务随 stream 释放而中止）， + // 写出循环入口的取消检查负责最终收尾 + if is_export_cancelled_now(&request.export_id) { + break; + } + } + } + + for (table_index, table_info) in tables.iter().enumerate() { // Check cancellation if is_export_cancelled(&request.export_id).await { on_progress(ExportProgress { @@ -1309,16 +1403,25 @@ pub async fn export_database_sql_core( .map_err(|e| format!("Failed to write file: {e}"))?; object_index += 1; } - match crate::schema::get_table_ddl_core( - state, - &request.connection_id, - &request.database, - &request.schema, - table_name, - None, - ) - .await + let ddl_result = match prefetched_table_metadata + .get_mut(table_index) + .and_then(|m| m.as_mut()) + .and_then(|m| m.ddl.take()) { + Some(result) => result, + None => { + crate::schema::get_table_ddl_core( + state, + &request.connection_id, + &request.database, + &request.schema, + table_name, + None, + ) + .await + } + }; + match ddl_result { Ok(ddl) => { let ddl = normalize_export_table_ddl(&ddl, Some(db_type)); writeln!(file, "{};\n", ddl).map_err(|e| format!("Failed to write file: {e}"))?; @@ -1336,15 +1439,24 @@ pub async fn export_database_sql_core( // Export data if request.include_data { // Get columns - let columns = match crate::schema::get_columns_core( - state, - &request.connection_id, - &request.database, - &request.schema, - table_name, - ) - .await + let columns_result = match prefetched_table_metadata + .get_mut(table_index) + .and_then(|m| m.as_mut()) + .and_then(|m| m.columns.take()) { + Some(result) => result, + None => { + crate::schema::get_columns_core( + state, + &request.connection_id, + &request.database, + &request.schema, + table_name, + ) + .await + } + }; + let columns = match columns_result { Ok(cols) => cols, Err(e) => { record_export_error( @@ -1688,6 +1800,7 @@ fn drop_table_if_exists_sql(table_name: &str, schema: &str, db_type: &DatabaseTy #[cfg(test)] mod tests { + use super::concurrent_metadata_prefetch_allowed; use super::{ build_database_sql_export, build_export_insert_statements, drop_table_if_exists_sql, filter_selected_table_infos, format_export_sql_literal, generate_postgres_sequence_create_ddl, @@ -1709,6 +1822,28 @@ mod tests { } } + #[test] + fn concurrent_prefetch_only_allowed_for_multi_connection_pools() { + use crate::connection::PoolKind; + use std::sync::Arc; + + // ChClient::new 只构造 HTTP 客户端，不发起连接 + let clickhouse = PoolKind::ClickHouse(crate::db::clickhouse_driver::ChClient::new( + "http://127.0.0.1:1", + None, + None, + std::time::Duration::from_secs(1), + )); + assert!(concurrent_metadata_prefetch_allowed(Some(&clickhouse))); + + // Agent（JDBC sidecar）请求超时覆盖排队时间，必须回退串行 + let agent = + PoolKind::Agent(Arc::new(tokio::sync::Mutex::new(crate::db::agent_driver::AgentDriverClient::test_stub()))); + assert!(!concurrent_metadata_prefetch_allowed(Some(&agent))); + + assert!(!concurrent_metadata_prefetch_allowed(None)); + } + #[test] fn filters_export_tables_by_selected_names() { let tables = vec![table("users", "TABLE"), table("orders", "TABLE"), table("active_users", "VIEW")]; diff --git a/crates/dbx-core/tests/database_export_prefetch.rs b/crates/dbx-core/tests/database_export_prefetch.rs new file mode 100644 index 0000000000..c2d79f8d86 --- /dev/null +++ b/crates/dbx-core/tests/database_export_prefetch.rs @@ -0,0 +1,210 @@ +//! 整库 SQL 导出的端到端回归测试（元数据并发预取后输出必须与逐表查询一致）。 +//! 需要本机 Docker；不可用时自动跳过。 + +use std::process::Command; + +use dbx_core::connection::AppState; +use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest}; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::storage::Storage; + +struct DockerPostgres { + name: String, + port: u16, +} + +impl Drop for DockerPostgres { + fn drop(&mut self) { + let _ = Command::new("docker").args(["rm", "-f", &self.name]).status(); + } +} + +fn docker_ready() -> bool { + Command::new("docker") + .args(["version", "--format", "{{.Server.Version}}"]) + .output() + .map(|output| output.status.success()) + .unwrap_or(false) +} + +fn start_docker_postgres() -> Option { + if !docker_ready() { + eprintln!("skipping docker-backed export test because Docker is unavailable"); + return None; + } + + let port = portpicker::pick_unused_port().expect("pick unused postgres port"); + let container = DockerPostgres { name: format!("dbx-export-prefetch-{}", uuid::Uuid::new_v4()), port }; + + let status = Command::new("docker") + .args([ + "run", + "-d", + "--rm", + "--name", + &container.name, + "-e", + "POSTGRES_PASSWORD=postgres", + "-e", + "POSTGRES_USER=postgres", + "-e", + "POSTGRES_DB=postgres", + "-p", + &format!("{port}:5432"), + "postgres:16-alpine", + ]) + .status() + .expect("start docker postgres"); + assert!(status.success(), "docker run postgres container should succeed"); + + // postgres 镜像 initdb 期间会短暂拉起再重启服务，单次探测可能命中引导期， + // 要求 SELECT 1 连续两次成功才算就绪 + let mut consecutive_ok = 0; + for _ in 0..120 { + let ready = Command::new("docker") + .args(["exec", &container.name, "psql", "-U", "postgres", "-d", "postgres", "-c", "SELECT 1"]) + .output() + .map(|output| output.status.success()) + .unwrap_or(false); + consecutive_ok = if ready { consecutive_ok + 1 } else { 0 }; + if consecutive_ok >= 2 { + return Some(container); + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + panic!("postgres container did not become ready"); +} + +fn psql(container: &DockerPostgres, sql: &str) { + let output = Command::new("docker") + .args(["exec", &container.name, "psql", "-U", "postgres", "-d", "postgres", "-v", "ON_ERROR_STOP=1", "-c", sql]) + .output() + .expect("run psql"); + assert!(output.status.success(), "psql failed: {}", String::from_utf8_lossy(&output.stderr)); +} + +fn postgres_test_config(id: &str, port: u16) -> ConnectionConfig { + ConnectionConfig { + id: id.to_string(), + name: id.to_string(), + db_type: DatabaseType::Postgres, + driver_profile: None, + driver_label: None, + url_params: None, + agent_java_options: Vec::new(), + host: "127.0.0.1".to_string(), + port, + username: "postgres".to_string(), + password: "postgres".to_string(), + database: Some("postgres".to_string()), + visible_databases: None, + visible_schemas: None, + attached_databases: Vec::new(), + init_script: None, + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 5, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 0, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: dbx_core::models::connection::default_redis_key_separator(), + redis_scan_page_size: None, + etcd_endpoints: String::new(), + gbase_server: String::new(), + informix_server: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + read_only: false, + is_production: false, + production_databases: vec![], + database_info: None, + } +} + +#[tokio::test] +async fn database_export_writes_structure_and_data_for_all_tables() { + let Some(container) = start_docker_postgres() else { + return; + }; + + psql( + &container, + "CREATE TABLE parent (id INT PRIMARY KEY, label TEXT NOT NULL);\ + CREATE TABLE child (id INT PRIMARY KEY, parent_id INT REFERENCES parent(id), note TEXT);\ + CREATE TABLE standalone_a (id INT PRIMARY KEY, payload JSONB);\ + CREATE TABLE standalone_b (id INT PRIMARY KEY, created_at TIMESTAMPTZ);\ + CREATE TABLE empty_table (id INT PRIMARY KEY);\ + INSERT INTO parent VALUES (1, 'alpha'), (2, 'beta');\ + INSERT INTO child VALUES (10, 1, 'first-child'), (11, 2, NULL);\ + INSERT INTO standalone_a VALUES (100, '{\"k\": \"v\"}');\ + INSERT INTO standalone_b VALUES (200, '2024-05-06T07:08:09Z');", + ); + + 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 connection_id = "export-prefetch-conn"; + state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); + + let file_path = dir.join("export.sql"); + let request = DatabaseExportRequest { + export_id: format!("export-prefetch-{}", uuid::Uuid::new_v4()), + connection_id: connection_id.to_string(), + database: "postgres".to_string(), + schema: "public".to_string(), + file_path: file_path.to_string_lossy().to_string(), + selected_tables: Vec::new(), + include_structure: true, + include_data: true, + include_objects: false, + drop_table_if_exists: true, + fail_on_error: true, + snapshot_session_id: None, + batch_size: 1000, + }; + + export_database_sql_core(&state, &request, |_progress| {}).await.expect("export should succeed"); + + let exported = std::fs::read_to_string(&file_path).expect("read exported file"); + + for table in ["parent", "child", "standalone_a", "standalone_b", "empty_table"] { + assert!( + exported.contains(&format!("CREATE TABLE \"public\".\"{table}\"")), + "exported SQL should contain CREATE TABLE for {table}:\n{exported}" + ); + } + assert!(!exported.contains("-- ERROR"), "exported SQL should not contain errors:\n{exported}"); + + for expected in ["'alpha'", "'beta'", "'first-child'", "\"k\"", "2024-05-06"] { + assert!(exported.contains(expected), "exported SQL should contain literal {expected}:\n{exported}"); + } + + // 依赖排序：被引用的 parent 必须先于 child 建表 + let parent_pos = exported.find("CREATE TABLE \"public\".\"parent\"").unwrap(); + let child_pos = exported.find("CREATE TABLE \"public\".\"child\"").unwrap(); + assert!(parent_pos < child_pos, "parent table DDL should precede child table DDL"); + + // 空表只导结构不导数据 + assert!( + !exported.contains("INSERT INTO \"public\".\"empty_table\""), + "empty table should not produce INSERT statements" + ); +}