feat(postgres): export all schemas in one SQL file

crates/dbx-core/src/database_export.rs

@@ -2586,10 +2586,21 @@ pub async fn export_database_sql_core(
     request: &DatabaseExportRequest,
     on_progress: impl Fn(ExportProgress) + Sync,
 ) -> Result<(), String> {
+    let db_type = state
+        .configs
+        .read()
+        .await
+        .get(&request.connection_id)
+        .map(|config| config.db_type)
+        .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?;
     // Keep the large export state machine on the heap. Besides making the
     // caller future small, this prevents the metadata-prefetch locals from
     // exhausting the bounded stack used by test and runtime worker threads.
-    let result = Box::pin(export_database_sql_core_inner(state, request, &on_progress)).await;
+    let result = if matches!(db_type, DatabaseType::Postgres) && request.schema.trim().is_empty() {
+        Box::pin(export_postgres_all_schemas_sql_core(state, request, &on_progress)).await
+    } else {
+        Box::pin(export_database_sql_core_inner(state, request, &on_progress)).await
+    };
     let metadata_session_id = database_export_client_session_id(&request.export_id);
     if let Err(error) =
         state.close_metadata_session_pool(&request.connection_id, Some(&request.database), &metadata_session_id).await
@@ -2611,6 +2622,185 @@ pub async fn export_database_sql_core(
     }
 }
 
+/// Destination directory that must exist (and stay on the same filesystem)
+/// before the export file may be written there, or `None` when the file path
+/// has no parent (e.g. a bare file name or the filesystem root). Shared by
+/// [`create_database_export_writer`] and the all-schemas export, which
+/// validates the destination up front instead of only after exporting every
+/// schema to temporary files.
+fn export_destination_parent_dir(file_path: &str) -> Option<&std::path::Path> {
+    let parent = std::path::Path::new(file_path).parent()?;
+    (!parent.as_os_str().is_empty()).then_some(parent)
+}
+
+async fn create_database_export_writer(
+    state: &Arc<crate::connection::AppState>,
+    request: &DatabaseExportRequest,
+) -> Result<DatabaseExportWriter, String> {
+    let mut expected_destination_identity = None;
+    if let Some(parent) = export_destination_parent_dir(&request.file_path) {
+        expected_destination_identity = ensure_export_destination_dir(state, parent).await?;
+    }
+    let file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to write file: {e}"))?;
+    // The directory check above and this `File::create` are separate
+    // operations: the mount can disappear and be replaced by something else
+    // at the same path in between. Re-check the identity of the handle we
+    // actually opened, not just the path, and refuse to keep a backup that
+    // landed on the wrong filesystem. See #6327.
+    let opened_destination_identity = export_destination_identity_for_file(&file);
+    if export_destination_identity_mismatch(
+        expected_destination_identity.as_ref(),
+        opened_destination_identity.as_ref(),
+    ) {
+        drop(file);
+        let _ = std::fs::remove_file(&request.file_path);
+        return Err(format!(
+            "Backup destination for {} changed while opening the output file -- the directory now \
+             resolves to a different filesystem than the one just verified. If a removable or network \
+             drive was disconnected and reconnected, retry the backup.",
+            request.file_path
+        ));
+    }
+    Ok(match request.output_compression {
+        DatabaseExportOutputCompression::None => DatabaseExportWriter::Plain(BufWriter::new(file)),
+        DatabaseExportOutputCompression::Gzip => {
+            DatabaseExportWriter::Gzip(Box::new(GzEncoder::new(BufWriter::new(file), Compression::default())))
+        }
+    })
+}
+
+fn postgres_export_schema_names(schemas: Vec<String>) -> Vec<String> {
+    let mut seen = HashSet::new();
+    schemas
+        .into_iter()
+        .map(|schema| schema.trim().to_string())
+        .filter(|schema| !schema.is_empty() && schema != "information_schema" && !schema.starts_with("pg_"))
+        .filter(|schema| seen.insert(schema.clone()))
+        .collect()
+}
+
+fn postgres_create_schema_sql(schema: &str) -> String {
+    format!("CREATE SCHEMA IF NOT EXISTS {};", quote_identifier(schema, &DatabaseType::Postgres))
+}
+
+async fn export_postgres_all_schemas_sql_core(
+    state: &Arc<crate::connection::AppState>,
+    request: &DatabaseExportRequest,
+    on_progress: impl Fn(ExportProgress) + Sync,
+) -> Result<(), String> {
+    let schemas = postgres_export_schema_names(
+        crate::schema::list_schemas_core(state, &request.connection_id, &request.database).await?,
+    );
+    if schemas.is_empty() {
+        return Err(format!("No exportable schemas found in database '{}'.", request.database));
+    }
+
+    // Validate the destination before exporting every schema to temporary
+    // files: the writer below only runs after the loop, which for large
+    // databases can be hours away, and a missing or unwritable destination
+    // must fail fast instead. This is an early fail, not a replacement -- the
+    // writer still performs its own checks when it opens the output file.
+    if let Some(parent) = export_destination_parent_dir(&request.file_path) {
+        ensure_export_destination_dir(state, parent).await?;
+    }
+
+    let temp_dir =
+        tempfile::tempdir().map_err(|error| format!("Failed to create temporary export directory: {error}"))?;
+    let result = async {
+        let mut schema_outputs = Vec::with_capacity(schemas.len());
+        let mut rows_exported = 0_u64;
+        let mut error_count = 0_u64;
+        let mut error_summary = None;
+
+        for (schema_index, schema_name) in schemas.iter().enumerate() {
+            let mut schema_request = request.clone();
+            schema_request.schema = schema_name.clone();
+            schema_request.file_path = temp_dir.path().join(format!("schema-{schema_index}.sql")).display().to_string();
+            schema_request.output_compression = DatabaseExportOutputCompression::None;
+
+            let terminal = Arc::new(std::sync::Mutex::new(None::<ExportProgress>));
+            let terminal_for_callback = terminal.clone();
+            let schema_name_for_callback = schema_name.clone();
+            let completed_rows = rows_exported;
+            let child_progress = |mut progress: ExportProgress| {
+                if matches!(progress.status, ExportStatus::Done | ExportStatus::Cancelled) {
+                    *terminal_for_callback.lock().expect("export terminal mutex poisoned") = Some(progress);
+                    return;
+                }
+                progress.export_id = request.export_id.clone();
+                progress.current_object = if progress.current_object.is_empty() {
+                    schema_name_for_callback.clone()
+                } else {
+                    format!("{schema_name_for_callback}: {}", progress.current_object)
+                };
+                progress.rows_exported = completed_rows.saturating_add(progress.rows_exported);
+                on_progress(progress);
+            };
+
+            Box::pin(export_database_sql_core_inner(state, &schema_request, child_progress)).await?;
+            let terminal = terminal.lock().expect("export terminal mutex poisoned").clone();
+            if terminal.as_ref().is_some_and(|progress| matches!(progress.status, ExportStatus::Cancelled)) {
+                return Err(EXPORT_CANCELLED_ERROR.to_string());
+            }
+            if let Some(progress) = terminal {
+                rows_exported = rows_exported.saturating_add(progress.rows_exported);
+                error_count = error_count.saturating_add(progress.error_count);
+                // Aggregate lenient failure summaries across schemas; keeping
+                // only the first schema's summary would understate the errors
+                // behind an error_count that accumulates over all schemas.
+                if let Some(summary) = progress.error_summary {
+                    error_summary = Some(match error_summary.take() {
+                        Some(existing) => format!("{existing}; {summary}"),
+                        None => summary,
+                    });
+                }
+            }
+            schema_outputs.push((schema_name.clone(), schema_request.file_path));
+        }
+
+        let mut file = create_database_export_writer(state, request).await?;
+        let timestamp = chrono::Local::now().format("%Y-%m-%d %H:%M:%S");
+        writeln!(file, "-- Database export: {}", request.database).map_err(|e| format!("Failed to write file: {e}"))?;
+        writeln!(file, "-- Date: {timestamp}").map_err(|e| format!("Failed to write file: {e}"))?;
+        writeln!(file, "-- Generated by DBX").map_err(|e| format!("Failed to write file: {e}"))?;
+        writeln!(file, "-- PostgreSQL schemas: {}", schemas.join(", "))
+            .map_err(|e| format!("Failed to write file: {e}"))?;
+        writeln!(file).map_err(|e| format!("Failed to write file: {e}"))?;
+        for schema_name in &schemas {
+            writeln!(file, "{}", postgres_create_schema_sql(schema_name))
+                .map_err(|e| format!("Failed to write file: {e}"))?;
+        }
+        writeln!(file).map_err(|e| format!("Failed to write file: {e}"))?;
+
+        for (schema_name, path) in schema_outputs {
+            writeln!(file, "-- Schema export: {schema_name}").map_err(|e| format!("Failed to write file: {e}"))?;
+            let mut source = std::io::BufReader::new(
+                std::fs::File::open(path).map_err(|e| format!("Failed to read temporary schema export: {e}"))?,
+            );
+            std::io::copy(&mut source, &mut file).map_err(|e| format!("Failed to combine schema export: {e}"))?;
+            writeln!(file).map_err(|e| format!("Failed to write file: {e}"))?;
+        }
+        file.finish()?;
+        on_progress(ExportProgress {
+            export_id: request.export_id.clone(),
+            current_object: request.database.clone(),
+            object_index: schemas.len(),
+            total_objects: schemas.len(),
+            rows_exported,
+            total_rows: None,
+            status: ExportStatus::Done,
+            error: None,
+            preparing: false,
+            error_count,
+            error_summary,
+        });
+        Ok(())
+    }
+    .await;
+
+    result
+}
+
 async fn export_database_sql_core_inner(
     state: &Arc<crate::connection::AppState>,
     request: &DatabaseExportRequest,
@@ -2658,38 +2848,7 @@ async fn export_database_sql_core_inner(
     ))
     .await?;
     // 4. Create file
-    let mut expected_destination_identity = None;
-    if let Some(parent) = std::path::Path::new(&request.file_path).parent() {
-        if !parent.as_os_str().is_empty() {
-            expected_destination_identity = ensure_export_destination_dir(state, parent).await?;
-        }
-    }
-    let file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to write file: {e}"))?;
-    // The directory check above and this `File::create` are separate
-    // operations: the mount can disappear and be replaced by something else
-    // at the same path in between. Re-check the identity of the handle we
-    // actually opened, not just the path, and refuse to keep a backup that
-    // landed on the wrong filesystem. See #6327.
-    let opened_destination_identity = export_destination_identity_for_file(&file);
-    if export_destination_identity_mismatch(
-        expected_destination_identity.as_ref(),
-        opened_destination_identity.as_ref(),
-    ) {
-        drop(file);
-        let _ = std::fs::remove_file(&request.file_path);
-        return Err(format!(
-            "Backup destination for {} changed while opening the output file -- the directory now \
-             resolves to a different filesystem than the one just verified. If a removable or network \
-             drive was disconnected and reconnected, retry the backup.",
-            request.file_path
-        ));
-    }
-    let mut file = match request.output_compression {
-        DatabaseExportOutputCompression::None => DatabaseExportWriter::Plain(BufWriter::new(file)),
-        DatabaseExportOutputCompression::Gzip => {
-            DatabaseExportWriter::Gzip(Box::new(GzEncoder::new(BufWriter::new(file), Compression::default())))
-        }
-    };
+    let mut file = create_database_export_writer(state, request).await?;
 
     let create_database_preamble = if request.include_create_database && matches!(db_type, DatabaseType::Mysql) {
         Some(mysql_database_export_preamble_for_request(state, request, &client_session_id).await)
@@ -3648,8 +3807,8 @@ mod tests {
     use super::{
         await_export_operation, await_export_stream_operation, clear_export_cancelled,
         concurrent_metadata_prefetch_allowed, database_export_metadata_prefetch_concurrency,
-        emit_database_export_cancelled, set_export_cancelled, snapshot_batch_cancelled, ExportStatus,
-        EXPORT_CANCELLED_ERROR,
+        emit_database_export_cancelled, postgres_create_schema_sql, postgres_export_schema_names, set_export_cancelled,
+        snapshot_batch_cancelled, ExportStatus, EXPORT_CANCELLED_ERROR,
     };
     use super::{
         build_database_export_object_source_sql, build_database_sql_export, build_export_insert_statements,
@@ -3948,6 +4107,26 @@ mod tests {
         assert!(!ddl.contains("VERSION"));
     }
 
+    #[test]
+    fn postgres_all_schema_export_filters_system_schemas_and_deduplicates() {
+        assert_eq!(
+            postgres_export_schema_names(vec![
+                " public ".to_string(),
+                "pg_catalog".to_string(),
+                "information_schema".to_string(),
+                "private".to_string(),
+                "private".to_string(),
+                "".to_string(),
+            ]),
+            vec!["public".to_string(), "private".to_string()]
+        );
+    }
+
+    #[test]
+    fn postgres_all_schema_export_creates_quoted_schema_if_missing() {
+        assert_eq!(postgres_create_schema_sql("tenant\"data"), "CREATE SCHEMA IF NOT EXISTS \"tenant\"\"data\";");
+    }
+
     #[test]
     fn postgres_extension_member_filter_keeps_user_overload_with_same_name() {
         let mut members = PostgresExtensionMembers::default();

crates/dbx-core/tests/live_postgres_all_schema_export.rs

@@ -0,0 +1,186 @@
+use dbx_core::connection::AppState;
+use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest, ExportStatus};
+use dbx_core::models::connection::{ConnectionConfig, DatabaseType};
+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, Mutex};
+use tokio_util::sync::CancellationToken;
+
+fn live_postgres_config(id: &str, database: &str) -> ConnectionConfig {
+    let host = std::env::var("DBX_LIVE_POSTGRES_EXPORT_HOST").expect("DBX_LIVE_POSTGRES_EXPORT_HOST");
+    let port =
+        std::env::var("DBX_LIVE_POSTGRES_EXPORT_PORT").ok().and_then(|value| value.parse::<u16>().ok()).unwrap_or(5432);
+    let username = std::env::var("DBX_LIVE_POSTGRES_EXPORT_USER").expect("DBX_LIVE_POSTGRES_EXPORT_USER");
+    let password = std::env::var("DBX_LIVE_POSTGRES_EXPORT_PASSWORD").expect("DBX_LIVE_POSTGRES_EXPORT_PASSWORD");
+
+    serde_json::from_value(serde_json::json!({
+        "id": id,
+        "name": id,
+        "db_type": DatabaseType::Postgres,
+        "host": host,
+        "port": port,
+        "username": username,
+        "password": password,
+        "database": database,
+        "ssl": false,
+        "connect_timeout_secs": 10,
+        "query_timeout_secs": 30,
+        "idle_timeout_secs": 60,
+        "keepalive_interval_secs": 0
+    }))
+    .expect("live PostgreSQL export config should deserialize")
+}
+
+#[tokio::test]
+#[ignore = "requires a live PostgreSQL account with CREATE DATABASE privilege"]
+async fn live_postgres_all_schema_export_restores_one_sql_file() {
+    let suffix = uuid::Uuid::new_v4().simple().to_string();
+    let source_database = format!("dbx_issue8336_src_{}", &suffix[..12]);
+    let target_database = format!("dbx_issue8336_dst_{}", &suffix[..12]);
+    let expected_export_database = source_database.clone();
+    let admin_connection_id = format!("postgres-all-schema-admin-{suffix}");
+    let source_connection_id = format!("postgres-all-schema-source-{suffix}");
+    let target_connection_id = format!("postgres-all-schema-target-{suffix}");
+    let dir = tempfile::tempdir().expect("create export temp directory");
+    let storage = Storage::open(&dir.path().join("storage.db")).await.expect("open temp storage");
+    let state = Arc::new(AppState::new(storage));
+    state
+        .configs
+        .write()
+        .await
+        .insert(admin_connection_id.clone(), live_postgres_config(&admin_connection_id, "postgres"));
+    for database in [&source_database, &target_database] {
+        Box::pin(execute_sql_statement(
+            &state,
+            &admin_connection_id,
+            "postgres",
+            &format!("CREATE DATABASE \"{database}\""),
+            None,
+            None,
+        ))
+        .await
+        .expect("create disposable PostgreSQL database");
+    }
+    state
+        .configs
+        .write()
+        .await
+        .insert(source_connection_id.clone(), live_postgres_config(&source_connection_id, &source_database));
+    state
+        .configs
+        .write()
+        .await
+        .insert(target_connection_id.clone(), live_postgres_config(&target_connection_id, &target_database));
+
+    for statement in [
+        "CREATE SCHEMA inventory",
+        "CREATE SCHEMA reporting",
+        "CREATE TABLE public.accounts (id integer PRIMARY KEY, name text NOT NULL)",
+        "CREATE TABLE inventory.products (id integer PRIMARY KEY, sku text NOT NULL)",
+        "CREATE TABLE reporting.daily_totals (day date PRIMARY KEY, amount numeric(12,2) NOT NULL)",
+        "CREATE VIEW reporting.positive_totals AS SELECT day, amount FROM reporting.daily_totals WHERE amount > 0",
+        "INSERT INTO public.accounts VALUES (1, 'Alice'), (2, 'Bob')",
+        "INSERT INTO inventory.products VALUES (10, 'SKU-10')",
+        "INSERT INTO reporting.daily_totals VALUES ('2026-09-07', 42.50)",
+    ] {
+        Box::pin(execute_sql_statement(&state, &source_connection_id, &source_database, statement, None, None))
+            .await
+            .expect("create source fixture");
+    }
+
+    let export_path = dir.path().join("postgres-all-schemas.sql");
+    let progress = Arc::new(Mutex::new(Vec::new()));
+    let progress_sink = progress.clone();
+    export_database_sql_core(
+        &state,
+        &DatabaseExportRequest {
+            export_id: format!("postgres-all-schema-export-{suffix}"),
+            connection_id: source_connection_id,
+            database: source_database.clone(),
+            schema: String::new(),
+            file_path: export_path.display().to_string(),
+            selected_tables: Vec::new(),
+            excluded_tables: Vec::new(),
+            include_structure: true,
+            include_data: true,
+            include_objects: true,
+            include_create_database: false,
+            drop_table_if_exists: false,
+            omit_auto_increment: false,
+            fail_on_error: true,
+            output_compression: Default::default(),
+            snapshot_session_id: None,
+            batch_size: 1000,
+        },
+        move |event| progress_sink.lock().expect("progress mutex poisoned").push(event),
+    )
+    .await
+    .expect("export every PostgreSQL schema");
+
+    let exported = std::fs::read_to_string(&export_path).expect("read combined export");
+    for expected in [
+        "CREATE SCHEMA IF NOT EXISTS \"public\";",
+        "CREATE SCHEMA IF NOT EXISTS \"inventory\";",
+        "CREATE SCHEMA IF NOT EXISTS \"reporting\";",
+        "CREATE TABLE \"public\".\"accounts\"",
+        "CREATE TABLE \"inventory\".\"products\"",
+        "CREATE TABLE \"reporting\".\"daily_totals\"",
+        "positive_totals",
+    ] {
+        assert!(exported.contains(expected), "missing {expected} in combined export:\n{exported}");
+    }
+    assert_eq!(exported.matches("-- Schema export:").count(), 3, "unexpected schema sections:\n{exported}");
+    assert!(
+        progress.lock().expect("progress mutex poisoned").iter().any(|event| {
+            matches!(event.status, ExportStatus::Done) && event.current_object == expected_export_database
+        }),
+        "combined export must emit a terminal Done event"
+    );
+
+    execute_sql_file_path(
+        &state,
+        &SqlFileRequest {
+            execution_id: format!("postgres-all-schema-import-{suffix}"),
+            connection_id: target_connection_id.clone(),
+            database: target_database.clone(),
+            file_path: export_path.display().to_string(),
+            continue_on_error: false,
+        },
+        &export_path,
+        CancellationToken::new(),
+        std::time::Instant::now(),
+        |_| {},
+    )
+    .await
+    .expect("restore combined PostgreSQL export");
+
+    let restored = Box::pin(execute_sql_statement(
+        &state,
+        &target_connection_id,
+        &target_database,
+        "SELECT
+           (SELECT count(*) FROM public.accounts),
+           (SELECT count(*) FROM inventory.products),
+           (SELECT amount::text FROM reporting.positive_totals WHERE day = DATE '2026-09-07')",
+        None,
+        None,
+    ))
+    .await
+    .expect("query restored schemas");
+    assert_eq!(restored.rows, vec![vec![serde_json::json!(2), serde_json::json!(1), serde_json::json!("42.50")]]);
+
+    for database in [&source_database, &target_database] {
+        Box::pin(execute_sql_statement(
+            &state,
+            &admin_connection_id,
+            "postgres",
+            &format!("DROP DATABASE \"{database}\" WITH (FORCE)"),
+            None,
+            None,
+        ))
+        .await
+        .expect("drop disposable PostgreSQL database");
+    }
+}