diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index b119e0ff9b..4594c6a88b 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1828,6 +1828,40 @@ pub async fn get_table_comment(pool: &Pool, schema: &str, table: &str) -> Result Ok(rows.first().and_then(|row| row.try_get::<_, Option>(0).ok().flatten()).filter(|s| !s.is_empty())) } +pub async fn get_table_partition_key(pool: &Pool, schema: &str, table: &str) -> Result, String> { + let schema = if schema.is_empty() { "public" } else { schema }; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + // Probe relkind first so PostgreSQL-compatible servers without the PG10 + // pg_get_partkeydef function keep ordinary-table DDL unchanged. + if postgres_query_cached(&client, postgres_partitioned_parent_sql(), &[&schema, &table]) + .await + .map_err(|e| e.to_string())? + .is_empty() + { + return Ok(None); + } + let rows = postgres_query_cached(&client, postgres_table_partition_key_sql(), &[&schema, &table]) + .await + .map_err(|e| e.to_string())?; + Ok(rows.first().and_then(|row| row.try_get::<_, Option>(0).ok().flatten()).filter(|s| !s.is_empty())) +} + +fn postgres_partitioned_parent_sql() -> &'static str { + "SELECT 1 \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind = 'p' \ + LIMIT 1" +} + +fn postgres_table_partition_key_sql() -> &'static str { + "SELECT pg_catalog.pg_get_partkeydef(c.oid) AS partition_key \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind = 'p' \ + LIMIT 1" +} + fn postgres_table_comment_sql() -> &'static str { "SELECT obj_description(c.oid) AS table_comment \ FROM pg_catalog.pg_class c \ @@ -4114,6 +4148,21 @@ mod tests { assert!(!sql.contains("ORDER BY")); } + #[test] + fn postgres_table_partition_key_sql_targets_partitioned_parents() { + let parent_sql = postgres_partitioned_parent_sql(); + let sql = postgres_table_partition_key_sql(); + + assert!(parent_sql.contains("n.nspname = $1")); + assert!(parent_sql.contains("c.relname = $2")); + assert!(parent_sql.contains("c.relkind = 'p'")); + assert!(sql.contains("pg_catalog.pg_get_partkeydef(c.oid)")); + assert!(sql.contains("n.nspname = $1")); + assert!(sql.contains("c.relname = $2")); + assert!(sql.contains("c.relkind = 'p'")); + assert!(!sql.contains("pg_get_expr(c.relpartbound")); + } + #[test] fn postgres_column_metadata_reads_identity_extra() { assert!(POSTGRES_COLUMNS_SQL.contains("a.attidentity")); @@ -4225,6 +4274,38 @@ mod tests { assert!(info.is_primary_key); } + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] + async fn postgres_partition_key_metadata_reads_parent_only() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); + let schema = format!("dbx_partition_meta_{}", std::process::id()); + let schema_ident = pg_quote_ident(&schema); + let parent = format!("{schema_ident}.parent"); + let child = format!("{schema_ident}.child"); + + execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); + let client = pool.get().await.expect("get postgres client"); + client + .batch_execute(&format!( + "CREATE TABLE {parent} (id integer, payload text) PARTITION BY RANGE (id); \ + CREATE TABLE {child} PARTITION OF {parent} FOR VALUES FROM (1) TO (10)" + )) + .await + .expect("create partitioned tables"); + + let parent_key = get_table_partition_key(&pool, &schema, "parent").await.expect("parent metadata"); + let child_key = get_table_partition_key(&pool, &schema, "child").await.expect("child metadata"); + let parent_ddl = crate::schema::pg_ddl(&pool, &schema, "parent").await.expect("parent ddl"); + + execute_query(&pool, &format!("DROP SCHEMA {schema_ident} CASCADE")).await.expect("drop schema"); + + assert_eq!(parent_key, Some("RANGE (id)".to_string())); + assert_eq!(child_key, None); + assert!(parent_ddl.contains(") PARTITION BY RANGE (id);"), "ddl: {parent_ddl}"); + assert!(!parent_ddl.contains("PARTITION OF")); + } + #[test] fn postgres_index_metadata_has_legacy_catalog_fallback() { assert!(POSTGRES_INDEXES_SQL.contains("ix.indnkeyatts")); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 0b8f5356f4..4503d3ec22 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -6277,6 +6277,45 @@ mod ddl_tests { assert!(ddl.contains("\"id\" integer generated by default as identity NOT NULL"), "ddl: {ddl}"); } + #[test] + fn postgres_table_ddl_includes_partition_key_for_parent_only() { + for partition_key in [ + "RANGE (created_at)", + "LIST (\"Tenant ID\")", + "HASH ((lower(code)))", + "RANGE (date_trunc('month'::text, created_at))", + ] { + let ddl = render_postgres_table_ddl_with_partition_key( + "public", + "events", + &[column("created_at", "timestamp without time zone")], + &[], + &[], + None, + Some(partition_key), + ); + + assert!(ddl.ends_with(&format!(") PARTITION BY {partition_key};\n")), "ddl: {ddl}"); + assert!(!ddl.contains("PARTITION OF")); + } + } + + #[test] + fn postgres_table_ddl_keeps_ordinary_table_unchanged() { + let ddl = render_postgres_table_ddl_with_partition_key( + "public", + "users", + &[column("id", "integer")], + &[], + &[], + None, + None, + ); + + assert!(ddl.ends_with(");\n"), "ddl: {ddl}"); + assert!(!ddl.contains("PARTITION BY")); + } + #[test] fn postgres_table_ddl_keeps_composite_foreign_key_together() { let columns = vec![column("a", "integer"), column("b", "integer"), column("c", "integer")]; @@ -6458,14 +6497,23 @@ pub fn opengauss_table_ddl_sql(schema: &str, table: &str) -> String { } pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { - let (columns, indexes, fkeys, table_comment) = tokio::try_join!( + let (columns, indexes, fkeys, table_comment, partition_key) = tokio::try_join!( db::postgres::get_columns(pool, schema, table), db::postgres::list_indexes(pool, schema, table), db::postgres::list_foreign_keys(pool, schema, table), async { db::postgres::get_table_comment(pool, schema, table).await }, + db::postgres::get_table_partition_key(pool, schema, table), )?; - Ok(render_postgres_table_ddl(schema, table, &columns, &indexes, &fkeys, table_comment.as_deref())) + Ok(render_postgres_table_ddl_with_partition_key( + schema, + table, + &columns, + &indexes, + &fkeys, + table_comment.as_deref(), + partition_key.as_deref(), + )) } pub fn render_postgres_table_ddl( @@ -6475,6 +6523,18 @@ pub fn render_postgres_table_ddl( indexes: &[db::IndexInfo], fkeys: &[db::ForeignKeyInfo], table_comment: Option<&str>, +) -> String { + render_postgres_table_ddl_with_partition_key(schema, table, columns, indexes, fkeys, table_comment, None) +} + +fn render_postgres_table_ddl_with_partition_key( + schema: &str, + table: &str, + columns: &[db::ColumnInfo], + indexes: &[db::IndexInfo], + fkeys: &[db::ForeignKeyInfo], + table_comment: Option<&str>, + partition_key: Option<&str>, ) -> String { let table_name = format!("{}.{}", pg_ident(schema), pg_ident(table)); let mut ddl = format!("CREATE TABLE {table_name} (\n"); @@ -6521,7 +6581,11 @@ pub fn render_postgres_table_ddl( ref_columns )); } - ddl.push_str("\n);\n"); + if let Some(partition_key) = partition_key.filter(|key| !key.trim().is_empty()) { + ddl.push_str(&format!("\n) PARTITION BY {partition_key};\n")); + } else { + ddl.push_str("\n);\n"); + } if let Some(comment) = table_comment.filter(|comment| !comment.trim().is_empty()) { ddl.push_str(&format!("\nCOMMENT ON TABLE {table_name} IS {};", sql_string(comment)));