fix(postgres): include partition key in table DDL

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<String>>(0).ok().flatten()).filter(|s| !s.is_empty()))
 }
 
+pub async fn get_table_partition_key(pool: &Pool, schema: &str, table: &str) -> Result<Option<String>, 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<String>>(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"));

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<String, String> {
-    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)));