diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index 54b3ac3eee..deb310bd0a 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -2133,7 +2133,7 @@ pub async fn export_database_sql_core( } let ddl = if request.include_structure { Some( - crate::schema::get_table_export_ddl_core( + crate::schema::get_table_relation_export_ddl_core( state, &request.connection_id, &request.database, @@ -2291,7 +2291,7 @@ pub async fn export_database_sql_core( { Some(result) => result, None => { - crate::schema::get_table_export_ddl_core( + crate::schema::get_table_relation_export_ddl_core( state, &request.connection_id, &request.database, diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 2cb3f5b5c5..a84a50e616 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -6498,9 +6498,10 @@ pub async fn list_owners_core( /// Whether to widen or normalize a single-table DDL fetch for its caller. /// /// Database export and table transfer render one relation at a time because -/// they already iterate every relation themselves. Interactive display paths -/// include PostgreSQL access statements and recurse through the partition -/// tree. Oracle export additionally requests portable DDL normalization. +/// they already iterate every relation themselves. A selected table structure +/// export and interactive display both recurse through the PostgreSQL +/// partition tree, while only display includes access statements. Oracle +/// exports additionally request portable DDL normalization. #[derive(Clone, Copy)] struct TableDdlOptions { include_postgres_access: bool, @@ -6511,7 +6512,9 @@ struct TableDdlOptions { impl TableDdlOptions { const SINGLE_RELATION: Self = Self { include_postgres_access: false, include_partitions: false, portable_oracle: false }; - const EXPORT: Self = Self { include_postgres_access: false, include_partitions: false, portable_oracle: true }; + const RELATION_EXPORT: Self = + Self { include_postgres_access: false, include_partitions: false, portable_oracle: true }; + const EXPORT: Self = Self { include_postgres_access: false, include_partitions: true, portable_oracle: true }; const DISPLAY: Self = Self { include_postgres_access: true, include_partitions: true, portable_oracle: false }; } @@ -6547,6 +6550,26 @@ pub async fn get_table_export_ddl_core( .await } +pub(crate) async fn get_table_relation_export_ddl_core( + state: &AppState, + connection_id: &str, + database: &str, + schema: &str, + table: &str, + object_type: Option, +) -> Result { + get_table_ddl_core_with_options( + state, + connection_id, + database, + schema, + table, + object_type, + TableDdlOptions::RELATION_EXPORT, + ) + .await +} + pub async fn get_table_display_ddl_core( state: &AppState, connection_id: &str, @@ -8520,6 +8543,21 @@ mod ddl_tests { } } + #[test] + fn table_structure_export_includes_partition_tree() { + assert!(TableDdlOptions::EXPORT.include_partitions); + assert!(TableDdlOptions::EXPORT.portable_oracle); + assert!(!TableDdlOptions::EXPORT.include_postgres_access); + + assert!(!TableDdlOptions::RELATION_EXPORT.include_partitions); + assert!(TableDdlOptions::RELATION_EXPORT.portable_oracle); + assert!(!TableDdlOptions::RELATION_EXPORT.include_postgres_access); + + assert!(TableDdlOptions::DISPLAY.include_partitions); + assert!(!TableDdlOptions::DISPLAY.portable_oracle); + assert!(TableDdlOptions::DISPLAY.include_postgres_access); + } + #[test] fn postgres_table_ddl_includes_column_comments() { let mut display_name = column("display_name", "text"); @@ -9566,8 +9604,8 @@ pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) - } /// Like `pg_ddl`, but for a partitioned table also emits `CREATE TABLE ... -/// PARTITION OF` for every existing partition, at any depth. Used only by the -/// interactive "view DDL" paths (`get_table_display_ddl_core`) — callers that +/// PARTITION OF` for every existing partition, at any depth. Used by selected +/// table structure exports and interactive "view DDL" paths — callers that /// iterate relations themselves must use `pg_ddl` instead (see its doc /// comment). Fetches the whole tree's metadata via a handful of batched, /// tree-wide queries (see `db::postgres::fetch_postgres_partition_tree` and diff --git a/crates/dbx-core/tests/database_export_partition_ddl.rs b/crates/dbx-core/tests/database_export_partition_ddl.rs index 5204206171..d01733b335 100644 --- a/crates/dbx-core/tests/database_export_partition_ddl.rs +++ b/crates/dbx-core/tests/database_export_partition_ddl.rs @@ -20,11 +20,14 @@ use dbx_core::storage::Storage; struct DockerPostgres { name: String, port: u16, + remove_on_drop: bool, } impl Drop for DockerPostgres { fn drop(&mut self) { - let _ = Command::new("docker").args(["rm", "-f", &self.name]).status(); + if self.remove_on_drop { + let _ = Command::new("docker").args(["rm", "-f", &self.name]).status(); + } } } @@ -37,13 +40,21 @@ fn docker_ready() -> bool { } fn start_docker_postgres(name_prefix: &str) -> Option { + if let Ok(name) = std::env::var("DBX_TEST_POSTGRES_CONTAINER") { + let port = std::env::var("DBX_TEST_POSTGRES_PORT") + .expect("DBX_TEST_POSTGRES_PORT must accompany DBX_TEST_POSTGRES_CONTAINER") + .parse() + .expect("DBX_TEST_POSTGRES_PORT must be a valid port"); + return Some(DockerPostgres { name, port, remove_on_drop: false }); + } if !docker_ready() { eprintln!("skipping docker-backed partition ddl test because Docker is unavailable"); return None; } let port = portpicker::pick_unused_port().expect("pick unused postgres port"); - let container = DockerPostgres { name: format!("{name_prefix}-{}", uuid::Uuid::new_v4()), port }; + let container = + DockerPostgres { name: format!("{name_prefix}-{}", uuid::Uuid::new_v4()), port, remove_on_drop: true }; let status = Command::new("docker") .args([ @@ -171,7 +182,10 @@ async fn database_export_of_partition_tree_has_no_duplicates_and_replays() { psql( &container, - "CREATE EXTENSION IF NOT EXISTS postgres_fdw;\ + "DROP TABLE IF EXISTS events CASCADE;\ + DROP TABLE IF EXISTS foreign_child_data CASCADE;\ + DROP SERVER IF EXISTS dbx_export_loopback CASCADE;\ + CREATE EXTENSION IF NOT EXISTS postgres_fdw;\ CREATE SERVER dbx_export_loopback FOREIGN DATA WRAPPER postgres_fdw \ OPTIONS (host 'localhost', dbname 'postgres', port '5432');\ CREATE USER MAPPING FOR CURRENT_USER SERVER dbx_export_loopback \ @@ -197,6 +211,45 @@ async fn database_export_of_partition_tree_has_no_duplicates_and_replays() { let connection_id = "export-partition-ddl-conn"; state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); + let selected_table_ddl = + dbx_core::schema::get_table_export_ddl_core(&state, connection_id, "postgres", "public", "events", None) + .await + .expect("selected table structure export should succeed"); + for table in ["events", "events_2026", "events_2027", "events_2027_h1", "events_2027_h2", "events_foreign"] { + let needle = format!("TABLE \"public\".\"{table}\""); + let count = selected_table_ddl.matches(&needle).count(); + assert_eq!(count, 1, "selected structure export should include {table} exactly once: {selected_table_ddl}"); + } + assert!( + selected_table_ddl.contains("CREATE FOREIGN TABLE \"public\".\"events_foreign\""), + "selected structure export should retain foreign partition syntax: {selected_table_ddl}" + ); + + let ordinary_table_ddl = dbx_core::schema::get_table_export_ddl_core( + &state, + connection_id, + "postgres", + "public", + "foreign_child_data", + None, + ) + .await + .expect("ordinary selected table structure export should succeed"); + assert_eq!(ordinary_table_ddl.matches("CREATE TABLE \"public\".\"foreign_child_data\"").count(), 1); + assert!(!ordinary_table_ddl.contains("PARTITION OF")); + + let missing_table_error = dbx_core::schema::get_table_export_ddl_core( + &state, + connection_id, + "postgres", + "public", + "missing_issue_6505_table", + None, + ) + .await + .expect_err("missing selected table should remain an error"); + assert!(missing_table_error.contains("was not found"), "unexpected missing-table error: {missing_table_error}"); + let file_path = dir.join("export.sql"); let request = DatabaseExportRequest { export_id: format!("export-partition-ddl-{}", uuid::Uuid::new_v4()), @@ -246,7 +299,7 @@ async fn database_export_of_partition_tree_has_no_duplicates_and_replays() { // The actual "restore" check: replay the exported SQL into a fresh // schema and confirm no `relation already exists` (or any other) error. - psql(&container, "CREATE SCHEMA replay"); + psql(&container, "DROP SCHEMA IF EXISTS replay CASCADE; CREATE SCHEMA replay"); let replayable = exported.replace("\"public\".", "\"replay\".").replace("SCHEMA \"public\"", "SCHEMA \"replay\""); let replay_file = dir.join("replay.sql"); std::fs::write(&replay_file, &replayable).unwrap();