feat(export): prefetch table metadata during database export

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<Mutex> 串行客户端）、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<Result<String, String>>,
+        columns: Option<Result<Vec<crate::db::ColumnInfo>, String>>,
+    }
+    const EXPORT_METADATA_PREFETCH_CONCURRENCY: usize = 8;
+    let mut prefetched_table_metadata: Vec<Option<PrefetchedTableMetadata>> = Vec::new();
+    prefetched_table_metadata.resize_with(tables.len(), || None);
+    // 防护门按「实际连接池种类」放行，而非数据库类型的能力标记：只有确认底层
+    // 支持并发请求的多连接池（Postgres/MySQL/ClickHouse）才启用并发预取。反例：
+    // 原生 SQL Server 是 Arc<Mutex<..>> 串行客户端；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")];

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<DockerPostgres> {
+    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"
+    );
+}