diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index ad6d1c9137..e2756c7cfb 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -109,6 +109,7 @@ pub enum PostgresColumnDefaultState { #[derive(Debug, Clone, Default, PartialEq, Eq)] pub struct PostgresTablePartitionLocalObjects { pub has_primary_key: bool, + pub unique_constraints: BTreeSet, pub foreign_keys: BTreeSet, pub indexes: BTreeSet, /// CHECK constraints with a local definition on this partition, as @@ -3662,6 +3663,9 @@ fn apply_partition_local_object_row( ) { match object_kind { "constraint" if object_type == "p" => entry.has_primary_key = true, + "constraint" if object_type == "u" && !object_name.is_empty() => { + entry.unique_constraints.insert(object_name); + } "constraint" if object_type == "f" && !object_name.is_empty() => { entry.foreign_keys.insert(object_name); } @@ -4078,7 +4082,7 @@ async fn list_indexes_for_relations_with_sql( key_is_expression, column_opclasses: key_opclasses, key_options, - constraint_backed: false, + constraint_backed: pg_row_try_bool(row, 13).unwrap_or(false), }); } Ok(result) @@ -4105,7 +4109,8 @@ fn postgres_indexes_for_relations_sql() -> &'static str { ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression, \ - array_agg(ix.indoption[(k.n - 1)::int] ORDER BY k.n) FILTER (WHERE k.n <= ix.indnkeyatts) AS key_options \ + array_agg(ix.indoption[(k.n - 1)::int] ORDER BY k.n) FILTER (WHERE k.n <= ix.indnkeyatts) AS key_options, \ + EXISTS (SELECT 1 FROM pg_constraint con WHERE con.conindid = i.oid) AS constraint_backed \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ @@ -4161,7 +4166,8 @@ fn postgres_indexes_for_relations_compat_sql() -> &'static str { AND a.attnum > 0 \ ORDER BY pos.n \ ) AS key_is_expression, \ - string_to_array(ix.indoption::text, ' ')::smallint[] AS key_options \ + string_to_array(ix.indoption::text, ' ')::smallint[] AS key_options, \ + EXISTS (SELECT 1 FROM pg_constraint con WHERE con.conindid = i.oid) AS constraint_backed \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ @@ -4399,7 +4405,7 @@ fn postgres_table_partition_local_objects_for_relations_query_tiers() -> [&'stat fn postgres_table_partition_local_objects_for_relations_sql() -> &'static str { "SELECT con.conrelid::bigint AS relid, 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ FROM pg_catalog.pg_constraint con \ - WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','f') \ + WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','u','f') \ AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ UNION ALL \ SELECT con.conrelid::bigint, 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ @@ -4438,7 +4444,7 @@ fn postgres_table_partition_local_objects_for_relations_sql() -> &'static str { fn postgres_table_partition_local_objects_for_relations_compat_sql() -> &'static str { "SELECT con.conrelid::bigint AS relid, 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ FROM pg_catalog.pg_constraint con \ - WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','f') \ + WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','u','f') \ UNION ALL \ SELECT con.conrelid::bigint, 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_constraint con \ @@ -4826,7 +4832,7 @@ fn postgres_table_partition_local_objects_sql() -> &'static str { FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ - WHERE n.nspname = $1 AND c.relname = $2 AND con.contype IN ('p','f') \ + WHERE n.nspname = $1 AND c.relname = $2 AND con.contype IN ('p','u','f') \ AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ UNION ALL \ SELECT 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ @@ -7687,7 +7693,8 @@ const POSTGRES_INDEXES_SQL: &str = "SELECT i.relname AS index_name, \ ix.indkey AS indkey, \ obj_description(i.oid, 'pg_class') AS index_comment, \ array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression, \ - array_agg(ix.indoption[(k.n - 1)::int] ORDER BY k.n) FILTER (WHERE k.n <= ix.indnkeyatts) AS key_options \ + array_agg(ix.indoption[(k.n - 1)::int] ORDER BY k.n) FILTER (WHERE k.n <= ix.indnkeyatts) AS key_options, \ + EXISTS (SELECT 1 FROM pg_constraint con WHERE con.conindid = i.oid) AS constraint_backed \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ @@ -7742,7 +7749,8 @@ const POSTGRES_INDEXES_COMPAT_SQL: &str = "SELECT i.relname AS index_name, \ AND a.attnum > 0 \ ORDER BY pos.n \ ) AS key_is_expression, \ - string_to_array(ix.indoption::text, ' ')::smallint[] AS key_options \ + string_to_array(ix.indoption::text, ' ')::smallint[] AS key_options, \ + EXISTS (SELECT 1 FROM pg_constraint con WHERE con.conindid = i.oid) AS constraint_backed \ FROM pg_index ix \ JOIN pg_class t ON t.oid = ix.indrelid \ JOIN pg_class i ON i.oid = ix.indexrelid \ @@ -7876,7 +7884,7 @@ async fn list_indexes_with_sql( key_is_expression, column_opclasses: key_opclasses, key_options, - constraint_backed: false, + constraint_backed: pg_row_try_bool(row, 12).unwrap_or(false), } }) .collect()) @@ -11669,7 +11677,7 @@ mod tests { assert!(info_compat_sql.contains("c.relkind IN ('r','p','f')")); assert!(info_compat_sql.contains("LIMIT 1")); assert!(local_objects_sql.contains("row_to_json(con)->>'conparentid'")); - assert!(local_objects_sql.contains("con.contype IN ('p','f')")); + assert!(local_objects_sql.contains("con.contype IN ('p','u','f')")); assert!(local_objects_sql.contains("i.inhrelid = idx.oid")); assert!(local_objects_sql.contains("con.contype = 'c' AND con.conislocal")); assert!(!local_objects_sql.contains("con.coninhcount = 0")); @@ -13046,6 +13054,8 @@ mod tests { postgres_indexes_for_relations_compat_sql(), ] { assert!(sql.contains("ix.indisunique AND ix.indisvalid")); + assert!(sql.contains("AS constraint_backed")); + assert!(sql.contains("con.conindid = i.oid")); } } diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 9879c1353a..add77487cb 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -10372,6 +10372,113 @@ mod ddl_tests { assert!(ddl.contains("\"id\" integer generated by default as identity NOT NULL"), "ddl: {ddl}"); } + #[test] + fn postgres_table_ddl_preserves_named_unique_and_primary_constraints() { + let mut id = column("id", "bigint"); + id.is_nullable = false; + id.is_primary_key = true; + let indexes = vec![ + db::IndexInfo { + name: "pk_accounts".to_string(), + columns: vec!["id".to_string()], + is_unique: true, + is_primary: true, + filter: None, + index_type: Some("btree".to_string()), + included_columns: None, + comment: None, + key_is_expression: Vec::new(), + column_opclasses: Vec::new(), + key_options: Vec::new(), + constraint_backed: true, + }, + db::IndexInfo { + name: "uq_accounts_code".to_string(), + columns: vec!["code".to_string()], + is_unique: true, + is_primary: false, + filter: None, + index_type: Some("btree".to_string()), + included_columns: None, + comment: None, + key_is_expression: Vec::new(), + column_opclasses: Vec::new(), + key_options: Vec::new(), + constraint_backed: true, + }, + db::IndexInfo { + name: "idx_accounts_display_name".to_string(), + columns: vec!["display_name".to_string()], + is_unique: true, + is_primary: false, + filter: None, + index_type: Some("btree".to_string()), + included_columns: None, + comment: None, + key_is_expression: Vec::new(), + column_opclasses: Vec::new(), + key_options: Vec::new(), + constraint_backed: false, + }, + ]; + let constraints = vec![ + db::ConstraintInfo { + name: "pk_accounts".to_string(), + constraint_type: "PRIMARY KEY".to_string(), + definition: "PRIMARY KEY (id)".to_string(), + columns: vec!["id".to_string()], + ref_schema: None, + ref_table: None, + ref_columns: Vec::new(), + match_type: None, + on_update: None, + on_delete: None, + deferrable: false, + initially_deferred: false, + enabled: true, + valid: true, + }, + db::ConstraintInfo { + name: "uq_accounts_code".to_string(), + constraint_type: "UNIQUE".to_string(), + definition: "UNIQUE (code) DEFERRABLE INITIALLY DEFERRED".to_string(), + columns: vec!["code".to_string()], + ref_schema: None, + ref_table: None, + ref_columns: Vec::new(), + match_type: None, + on_update: None, + on_delete: None, + deferrable: true, + initially_deferred: true, + enabled: true, + valid: true, + }, + ]; + + let ddl = render_postgres_table_ddl_with_constraints_and_partition_info( + "public", + "accounts", + &[id], + &indexes, + &[], + &constraints, + &[], + None, + &db::postgres::PostgresTablePartitionInfo::default(), + &db::postgres::PostgresTablePartitionLocalObjects::default(), + ); + + assert!(ddl.contains("CONSTRAINT \"pk_accounts\" PRIMARY KEY (id)"), "ddl: {ddl}"); + assert!( + ddl.contains("CONSTRAINT \"uq_accounts_code\" UNIQUE (code) DEFERRABLE INITIALLY DEFERRED"), + "ddl: {ddl}" + ); + assert!(!ddl.contains("CREATE UNIQUE INDEX \"pk_accounts\""), "ddl: {ddl}"); + assert!(!ddl.contains("CREATE UNIQUE INDEX \"uq_accounts_code\""), "ddl: {ddl}"); + assert!(ddl.contains("CREATE UNIQUE INDEX \"idx_accounts_display_name\""), "ddl: {ddl}"); + } + #[test] fn postgres_table_ddl_renders_owned_serial_markers_without_external_defaults() { for (column_name, data_type, serial_type) in [ @@ -10545,6 +10652,68 @@ mod ddl_tests { assert!(!ddl.contains("\"payload\" text")); } + #[test] + fn postgres_partition_ddl_preserves_local_unique_without_local_primary_key() { + let indexes = vec![db::IndexInfo { + name: "events_2026_code_key".to_string(), + columns: vec!["code".to_string()], + is_unique: true, + is_primary: false, + filter: None, + index_type: Some("btree".to_string()), + included_columns: None, + comment: None, + key_is_expression: Vec::new(), + column_opclasses: Vec::new(), + key_options: Vec::new(), + constraint_backed: true, + }]; + let constraints = vec![db::ConstraintInfo { + name: "events_2026_code_key".to_string(), + constraint_type: "UNIQUE".to_string(), + definition: "UNIQUE (code)".to_string(), + columns: vec!["code".to_string()], + ref_schema: None, + ref_table: None, + ref_columns: Vec::new(), + match_type: None, + on_update: None, + on_delete: None, + deferrable: false, + initially_deferred: false, + enabled: true, + valid: true, + }]; + let partition_info = db::postgres::PostgresTablePartitionInfo { + is_partition: true, + parent_schema: Some("public".to_string()), + parent_table: Some("events".to_string()), + bound: Some("DEFAULT".to_string()), + ..Default::default() + }; + let partition_local_objects = db::postgres::PostgresTablePartitionLocalObjects { + unique_constraints: BTreeSet::from(["events_2026_code_key".to_string()]), + indexes: BTreeSet::from(["events_2026_code_key".to_string()]), + ..Default::default() + }; + + let ddl = render_postgres_table_ddl_with_constraints_and_partition_info( + "public", + "events_2026", + &[column("code", "text")], + &indexes, + &[], + &constraints, + &[], + None, + &partition_info, + &partition_local_objects, + ); + + assert!(ddl.contains("CONSTRAINT \"events_2026_code_key\" UNIQUE (code)"), "ddl: {ddl}"); + assert!(!ddl.contains("CREATE UNIQUE INDEX"), "ddl: {ddl}"); + } + #[test] fn postgres_partition_ddl_skips_inherited_constraints_and_indexes() { let mut id = column("id", "integer"); @@ -11394,15 +11563,17 @@ pub async fn sqlite_ddl(pool: &db::sqlite::SqliteHandle, schema: &str, table: &s /// duplicate every partition's `CREATE TABLE` (once from the parent's DDL, /// once from the caller's own loop over that same child relation). pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { - let (columns, indexes, fkeys, table_comment, partition_info, trigger_definitions, check_constraints) = 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_info(pool, schema, table), - db::postgres::list_trigger_definitions(pool, schema, table), - db::postgres::list_check_constraints(pool, schema, table), - )?; + let (columns, indexes, fkeys, constraints, table_comment, partition_info, trigger_definitions, check_constraints) = + 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), + db::postgres::list_constraints(pool, schema, table), + async { db::postgres::get_table_comment(pool, schema, table).await }, + db::postgres::get_table_partition_info(pool, schema, table), + db::postgres::list_trigger_definitions(pool, schema, table), + db::postgres::list_check_constraints(pool, schema, table), + )?; let partition_local_objects = if partition_info.is_partition { db::postgres::get_table_partition_local_objects(pool, schema, table).await? } else { @@ -11410,12 +11581,13 @@ pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) - }; Ok(append_postgres_trigger_definitions( - render_postgres_table_ddl_with_partition_info( + render_postgres_table_ddl_with_constraints_and_partition_info( schema, table, &columns, &indexes, &fkeys, + &constraints, &check_constraints, table_comment.as_deref(), &partition_info, @@ -12169,6 +12341,32 @@ fn render_postgres_table_ddl_with_partition_info( table_comment: Option<&str>, partition_info: &db::postgres::PostgresTablePartitionInfo, partition_local_objects: &db::postgres::PostgresTablePartitionLocalObjects, +) -> String { + render_postgres_table_ddl_with_constraints_and_partition_info( + schema, + table, + columns, + indexes, + fkeys, + &[], + check_constraints, + table_comment, + partition_info, + partition_local_objects, + ) +} + +fn render_postgres_table_ddl_with_constraints_and_partition_info( + schema: &str, + table: &str, + columns: &[db::ColumnInfo], + indexes: &[db::IndexInfo], + fkeys: &[db::ForeignKeyInfo], + constraints: &[db::ConstraintInfo], + check_constraints: &[(String, String)], + table_comment: Option<&str>, + partition_info: &db::postgres::PostgresTablePartitionInfo, + partition_local_objects: &db::postgres::PostgresTablePartitionLocalObjects, ) -> String { let table_name = format!("{}.{}", pg_ident(schema), pg_ident(table)); let partition_parent = partition_info @@ -12216,14 +12414,38 @@ fn render_postgres_table_ddl_with_partition_info( .collect::>() }; - let pks: Vec<&str> = if !is_partition || partition_local_objects.has_primary_key { - columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.as_str()).collect() - } else { - Vec::new() - }; - if !pks.is_empty() { - definition_lines - .push(format!(" PRIMARY KEY ({})", pks.iter().map(|key| pg_ident(key)).collect::>().join(", "))); + let primary_constraints = constraints + .iter() + .filter(|constraint| constraint.constraint_type == "PRIMARY KEY" && !constraint.definition.trim().is_empty()) + .collect::>(); + let unique_constraints = constraints + .iter() + .filter(|constraint| constraint.constraint_type == "UNIQUE" && !constraint.definition.trim().is_empty()) + .collect::>(); + if !is_partition || partition_local_objects.has_primary_key { + if primary_constraints.is_empty() { + let pks: Vec<&str> = columns.iter().filter(|c| c.is_primary_key).map(|c| c.name.as_str()).collect(); + if !pks.is_empty() { + definition_lines.push(format!( + " PRIMARY KEY ({})", + pks.iter().map(|key| pg_ident(key)).collect::>().join(", ") + )); + } + } else { + for constraint in &primary_constraints { + definition_lines.push(format!( + " CONSTRAINT {} {}", + pg_ident(&constraint.name), + constraint.definition.trim() + )); + } + } + } + for constraint in unique_constraints { + if is_partition && !partition_local_objects.unique_constraints.contains(&constraint.name) { + continue; + } + definition_lines.push(format!(" CONSTRAINT {} {}", pg_ident(&constraint.name), constraint.definition.trim())); } for fk_group in group_foreign_keys_by_name(fkeys) { let Some(first_fk) = fk_group.first() else { @@ -12343,7 +12565,7 @@ fn render_postgres_table_ddl_with_partition_info( } for idx in indexes { - if idx.is_primary { + if idx.is_primary || idx.constraint_backed { continue; } if is_partition && !partition_local_objects.indexes.contains(&idx.name) {