diff --git a/src/build/mod.rs b/src/build/mod.rs index 2b05494..a3f20c5 100644 --- a/src/build/mod.rs +++ b/src/build/mod.rs @@ -203,6 +203,23 @@ impl Builder { defn: Vec, drop_stmt: Vec, no_owner: bool, + ) -> Result<(), String> { + self.add_item_with_comment_target( + item, defn, drop_stmt, no_owner, None, + ) + } + + /// Like [`add_item`], but lets the caller supply the exact object + /// reference (e.g. with an `(argtypes)` signature) used in the + /// item's `COMMENT ON` statement, for object types whose comment + /// target isn't just `namespace.tag`. + fn add_item_with_comment_target( + &mut self, + item: &Item, + defn: Vec, + drop_stmt: Vec, + no_owner: bool, + comment_target: Option, ) -> Result<(), String> { let namespace = item.definition.schema().unwrap_or_default().to_string(); @@ -232,12 +249,17 @@ impl Builder { &owner, dump_id, comment, + comment_target, )?; } Ok(()) } - /// Add a COMMENT ON entry tied to its parent (ports _add_comment) + /// Add a COMMENT ON entry tied to its parent (ports _add_comment). + /// `target`, when given, is used verbatim as the object reference + /// (already fully quoted/qualified); otherwise it is built from + /// `namespace` and `tag`. + #[allow(clippy::too_many_arguments)] fn add_comment( &mut self, desc: &str, @@ -246,16 +268,27 @@ impl Builder { owner: &str, parent_dump_id: i32, comment: &str, + target: Option, ) -> Result<(), String> { // extensions record the schema they install into as their - // namespace, but COMMENT ON EXTENSION takes an unqualified name - let name = if desc == "EXTENSION" { - quote_ident(tag) - } else if namespace.is_empty() { - tag.to_string() - } else { - format!("{namespace}.{tag}") - }; + // namespace, but COMMENT ON EXTENSION takes an unqualified name. + // FUNCTION tags already carry their own (argtypes) signature, + // so they pass through unquoted rather than as a plain ident. + let name = target.unwrap_or_else(|| { + if desc == "EXTENSION" { + quote_ident(tag) + } else if desc == "FUNCTION" { + if namespace.is_empty() { + tag.to_string() + } else { + format!("{}.{}", quote_ident(namespace), tag) + } + } else if namespace.is_empty() { + quote_ident(tag) + } else { + format!("{}.{}", quote_ident(namespace), quote_ident(tag)) + } + }); let defn = vec![ String::from("COMMENT"), String::from("ON"), @@ -372,12 +405,21 @@ impl Builder { options.push("HYPOTHETICAL".into()); } create.push(format!("({})", options.join(", "))); + let signature = format!("({})", args.join(", ")); let drop = vec![ "DROP AGGREGATE IF EXISTS".into(), self.item_name(item), - format!("({})", args.join(", ")), + signature.clone(), ]; - self.add_item(item, create, drop, false) + // COMMENT ON AGGREGATE needs the argument signature appended + let comment_target = format!("{} {signature}", self.item_name(item)); + self.add_item_with_comment_target( + item, + create, + drop, + false, + Some(comment_target), + ) } fn dump_cast(&mut self, item: &Item) -> Result<(), String> { @@ -407,8 +449,14 @@ impl Builder { if d.implicit == Some(true) { create.push("AS IMPLICIT".into()); } - let drop = vec!["DROP CAST IF EXISTS".into(), name]; - self.add_item(item, create, drop, false) + let drop = vec!["DROP CAST IF EXISTS".into(), name.clone()]; + self.add_item_with_comment_target( + item, + create, + drop, + false, + Some(name), + ) } fn dump_collation(&mut self, item: &Item) -> Result<(), String> { @@ -857,16 +905,25 @@ impl Builder { options.push("MERGES".into()); } create.push(format!("({})", options.join(", "))); + let signature = format!( + "({}, {})", + d.left_arg.as_deref().unwrap_or("NONE"), + d.right_arg.as_deref().unwrap_or("NONE") + ); let drop = vec![ "DROP OPERATOR IF EXISTS".into(), - name, - format!( - "({}, {})", - d.left_arg.as_deref().unwrap_or("NONE"), - d.right_arg.as_deref().unwrap_or("NONE") - ), + name.clone(), + signature.clone(), ]; - self.add_item(item, create, drop, false) + // COMMENT ON OPERATOR needs the argument signature appended + let comment_target = format!("{name} {signature}"); + self.add_item_with_comment_target( + item, + create, + drop, + false, + Some(comment_target), + ) } fn dump_publication(&mut self, item: &Item) -> Result<(), String> { @@ -1099,7 +1156,7 @@ impl Builder { if value { "INCLUDING" } else { "EXCLUDING" } .into(), ); - create.push(format!("include_{field}")); + create.push(field.to_uppercase()); } } } else { @@ -1118,9 +1175,16 @@ impl Builder { for fk in d.foreign_keys.as_deref().unwrap_or_default() { inner.push(render_foreign_key(fk)); } + for check in d.check_constraints.as_deref().unwrap_or_default() + { + inner.push(format!( + "CONSTRAINT {} CHECK ({})", + check.name, check.expression + )); + } create.push(inner.join(", ")); - create.push(")".into()); } + create.push(")".into()); if let Some(parents) = &d.parents { create.push("INHERITS".into()); create.push(format!("({})", parents.join(", "))); @@ -1157,6 +1221,26 @@ impl Builder { vec!["DROP TABLE IF EXISTS".into(), self.item_name(item)]; self.add_item(item, create, drop, false)?; } + let dump_id = self.dump_id_map[&item.id]; + for column in d.columns.as_deref().unwrap_or_default() { + if let Some(comment) = &column.comment { + let target = format!( + "{}.{}.{}", + quote_ident(&d.schema), + quote_ident(&d.name), + quote_ident(&column.name) + ); + self.add_comment( + "COLUMN", + &d.schema, + &format!("{}.{}", d.name, column.name), + &d.owner, + dump_id, + comment, + Some(target), + )?; + } + } for index in d.indexes.as_deref().unwrap_or_default() { self.dump_index(index, item, &d.schema, &d.owner)?; } @@ -1217,6 +1301,7 @@ impl Builder { &table.owner, dump_id, comment, + None, )?; } Ok(()) @@ -1255,6 +1340,7 @@ impl Builder { owner, dump_id, comment, + None, )?; } Ok(()) @@ -1280,6 +1366,13 @@ impl Builder { None, )?; if let Some(comment) = &trigger.comment { + // a trigger name isn't schema-qualifiable; COMMENT ON + // TRIGGER requires the owning table instead + let target = format!( + "{} ON {}", + quote_ident(&name), + self.item_name(parent) + ); self.add_comment( "TRIGGER", &table.schema, @@ -1287,6 +1380,7 @@ impl Builder { &table.owner, dump_id, comment, + Some(target), )?; } Ok(()) @@ -1504,6 +1598,7 @@ impl Builder { &self.superuser.clone(), dump_id, comment, + None, )?; } Ok(()) @@ -1729,11 +1824,15 @@ impl Builder { columns.iter().map(view_column_name).collect(); create.push(format!("({})", names.join(", "))); } + let mut with_options = Vec::new(); if let Some(check_option) = &d.check_option { - create.push(format!("WITH (check_option = {check_option})")); + with_options.push(format!("check_option = {check_option}")); } if d.security_barrier == Some(true) { - create.push("WITH (security_barrier = true)".into()); + with_options.push("security_barrier = true".into()); + } + if !with_options.is_empty() { + create.push(format!("WITH ({})", with_options.join(", "))); } create.push("AS".into()); create.push(d.query.clone().unwrap_or_default()); @@ -2097,7 +2196,7 @@ mod tests { use super::*; use crate::constants::ObjectType; - use crate::models::{Table, Trigger}; + use crate::models::{CheckConstraint, LikeTable, Table, Trigger, View}; fn column(name: &str, data_type: &str, not_null: bool) -> Column { Column { @@ -2227,4 +2326,195 @@ mod tests { EXECUTE FUNCTION test.emit()" ); } + + fn base_table() -> Table { + Table { + name: "orders".into(), + schema: "public".into(), + owner: "app".into(), + sql: None, + unlogged: None, + from_type: None, + parents: None, + like_table: None, + columns: None, + indexes: None, + primary_key: None, + check_constraints: None, + unique_constraints: None, + foreign_keys: None, + triggers: None, + partition: None, + partitions: None, + access_method: None, + storage_parameters: None, + tablespace: None, + index_tablespace: None, + server: None, + options: None, + comment: None, + } + } + + /// Render `table` and return the TABLE entry's definition SQL + fn table_defn(table: Table) -> String { + let item = Item { + id: 1, + desc: ObjectType::Table, + definition: Definition::Table(table), + dependencies: BTreeSet::new(), + }; + let dump = libpgdump::new("t", "UTF-8", "18.0").unwrap(); + let mut builder = Builder { + dump, + dump_id_map: HashMap::new(), + superuser: "postgres".into(), + }; + builder.dump_item(&item).unwrap(); + builder + .dump + .entries() + .iter() + .find(|e| e.desc == libpgdump::ObjectType::Table) + .and_then(|e| e.defn.clone()) + .expect("a TABLE entry") + } + + #[test] + fn renders_like_table_include_options() { + let mut table = base_table(); + table.like_table = Some(LikeTable { + name: "public.template".into(), + include_comments: Some(true), + include_constraints: None, + include_defaults: None, + include_generated: None, + include_identity: None, + include_indexes: Some(false), + include_statistics: None, + include_storage: None, + include_all: None, + }); + assert_eq!( + table_defn(table), + "CREATE TABLE public.orders ( LIKE public.template \ + INCLUDING COMMENTS EXCLUDING INDEXES );\n" + ); + } + + #[test] + fn renders_table_check_constraint() { + let mut table = base_table(); + table.columns = Some(vec![column("qty", "integer", true)]); + table.check_constraints = Some(vec![CheckConstraint { + name: "ck_positive".into(), + expression: "qty > 0".into(), + }]); + assert_eq!( + table_defn(table), + "CREATE TABLE public.orders ( qty integer NOT NULL, \ + CONSTRAINT ck_positive CHECK (qty > 0) );\n" + ); + } + + /// Render `item` and return every COMMENT entry's definition SQL + fn comment_defns(item: &Item) -> Vec { + let dump = libpgdump::new("t", "UTF-8", "18.0").unwrap(); + let mut builder = Builder { + dump, + dump_id_map: HashMap::new(), + superuser: "postgres".into(), + }; + builder.dump_item(item).unwrap(); + builder + .dump + .entries() + .iter() + .filter(|e| e.desc == libpgdump::ObjectType::Comment) + .filter_map(|e| e.defn.clone()) + .collect() + } + + #[test] + fn renders_trigger_comment_on_owning_table() { + let mut table = base_table(); + table.columns = Some(vec![column("id", "integer", true)]); + let mut trigger = constraint_trigger(false); + trigger.comment = Some("audits inserts".into()); + table.triggers = Some(vec![trigger]); + let item = Item { + id: 1, + desc: ObjectType::Table, + definition: Definition::Table(table), + dependencies: BTreeSet::new(), + }; + assert_eq!( + comment_defns(&item), + vec![ + "COMMENT ON TRIGGER emit ON public.orders IS $$audits \ + inserts$$;\n;\n" + ] + ); + } + + #[test] + fn renders_column_comment() { + let mut table = base_table(); + let mut qty = column("qty", "integer", true); + qty.comment = Some("quantity ordered".into()); + table.columns = Some(vec![qty]); + let item = Item { + id: 1, + desc: ObjectType::Table, + definition: Definition::Table(table), + dependencies: BTreeSet::new(), + }; + assert_eq!( + comment_defns(&item), + vec![ + "COMMENT ON COLUMN public.orders.qty IS $$quantity \ + ordered$$;\n;\n" + ] + ); + } + + #[test] + fn renders_view_with_check_option_and_security_barrier() { + let item = Item { + id: 1, + desc: ObjectType::View, + definition: Definition::View(View { + name: "active_orders".into(), + schema: "public".into(), + owner: "app".into(), + sql: None, + recursive: None, + columns: None, + check_option: Some("local".into()), + security_barrier: Some(true), + query: Some("SELECT 1".into()), + comment: None, + }), + dependencies: BTreeSet::new(), + }; + let dump = libpgdump::new("t", "UTF-8", "18.0").unwrap(); + let mut builder = Builder { + dump, + dump_id_map: HashMap::new(), + superuser: "postgres".into(), + }; + builder.dump_item(&item).unwrap(); + let defn = builder + .dump + .entries() + .iter() + .find(|e| e.desc == libpgdump::ObjectType::View) + .and_then(|e| e.defn.clone()) + .expect("a VIEW entry"); + assert_eq!( + defn, + "CREATE VIEW public.active_orders WITH (check_option = \ + local, security_barrier = true) AS SELECT 1;\n" + ); + } } diff --git a/src/ddl/function.rs b/src/ddl/function.rs index 351b3f7..01a9309 100644 --- a/src/ddl/function.rs +++ b/src/ddl/function.rs @@ -64,11 +64,32 @@ pub(crate) fn create_function( .child_of_kind("NonReservedWord_or_Sconst") .map(|n| unquote(n.text(src))); } else if option.has("kw_as") { - if let Some(body) = option.find("Sconst") { - function.definition = Some(string_value(&body, src)); + if let Some(func_as) = option.child_of_kind("func_as") { + let sconsts = func_as.find_all("Sconst"); + match sconsts.as_slice() { + [object_file, link_symbol] => { + function.object_file = + Some(string_value(object_file, src)); + function.link_symbol = + Some(string_value(link_symbol, src)); + } + [body] => { + function.definition = Some(string_value(body, src)); + } + _ => {} + } } } else if option.has("kw_window") { function.window = Some(true); + } else if option.has("kw_transform") { + let transform_types: Vec = option + .find_all("Typename") + .iter() + .map(|n| n.text(src).to_string()) + .collect(); + if !transform_types.is_empty() { + function.transform_types = Some(transform_types); + } } else if let Some(common) = option.child_of_kind("common_func_opt_item") { @@ -264,6 +285,32 @@ mod tests { ); } + #[test] + fn parses_c_language_function() { + let Statement::CreateFunction(function) = parse_one( + "CREATE FUNCTION test.c_fn(a integer) RETURNS integer \ + LANGUAGE c AS 'ext_module', 'c_fn';", + ) else { + panic!("expected CreateFunction") + }; + assert_eq!(function.language, Some("c".into())); + assert_eq!(function.object_file, Some("ext_module".into())); + assert_eq!(function.link_symbol, Some("c_fn".into())); + assert_eq!(function.definition, None); + } + + #[test] + fn parses_transform_types() { + let Statement::CreateFunction(function) = parse_one( + "CREATE FUNCTION test.fn(a hstore) RETURNS integer \ + LANGUAGE plpython3u TRANSFORM FOR TYPE hstore \ + AS $$ return 1 $$;", + ) else { + panic!("expected CreateFunction") + }; + assert_eq!(function.transform_types, Some(vec!["hstore".into()])); + } + #[test] fn parses_unqualified_function() { let Statement::CreateFunction(function) = parse_one( diff --git a/src/ddl/mod.rs b/src/ddl/mod.rs index fe1401e..4ba488a 100644 --- a/src/ddl/mod.rs +++ b/src/ddl/mod.rs @@ -24,6 +24,13 @@ use crate::models; #[derive(Clone, Debug, PartialEq)] pub enum Statement { CreateTable(Box), + /// CREATE TABLE ... PARTITION OF `parent` FOR VALUES ... — folded + /// into the parent table's `partitions` list during pull assembly + /// rather than modeled as a standalone table + CreateTablePartition { + parent: QualifiedName, + partition: models::TablePartition, + }, /// CREATE INDEX — `table` is the qualified relation name CreateIndex { table: QualifiedName, diff --git a/src/ddl/object.rs b/src/ddl/object.rs index d66c621..063e8b2 100644 --- a/src/ddl/object.rs +++ b/src/ddl/object.rs @@ -415,6 +415,48 @@ pub(crate) fn unstring(text: &str) -> String { text.to_string() } +/// Parse a `reloptions` node (`key = value, ...` inside `WITH (...)`) +/// into a string map; values keep their raw text (dequoted if a +/// string literal) so they round-trip unchanged through +/// `utils::raw_value` +pub(crate) fn reloptions( + node: &Node, + src: &str, +) -> Option> { + let elems = node.find_all("reloption_elem"); + if elems.is_empty() { + return None; + } + let mut map = serde_json::Map::new(); + for elem in elems { + let labels = elem.find_all("ColLabel"); + let Some(first) = labels.first() else { + continue; + }; + let key = match labels.get(1) { + Some(second) => { + format!( + "{}.{}", + unquote(first.text(src)), + unquote(second.text(src)) + ) + } + None => unquote(first.text(src)), + }; + // A bare option (`WITH (security_barrier)`) has no def_arg; + // PostgreSQL treats it as shorthand for `= true`. + let value = match elem.child_of_kind("def_arg") { + Some(value) => value + .child_of_kind("Sconst") + .map(|s| string_value(&s, src)) + .unwrap_or_else(|| value.text(src).to_string()), + None => "true".to_string(), + }; + map.insert(key, serde_json::Value::String(value)); + } + (!map.is_empty()).then_some(map) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/ddl/table.rs b/src/ddl/table.rs index 1c9e6f6..40d2721 100644 --- a/src/ddl/table.rs +++ b/src/ddl/table.rs @@ -1,32 +1,92 @@ //! CreateStmt / IndexStmt / AlterTableStmt → table models +use serde_json::Value; use tree_sitter::Node; +use crate::ddl::object::{reloptions, string_value}; use crate::ddl::{ - NodeExt, Statement, TableConstraint, qualified_name, unquote, + NodeExt, Statement, TableConstraint, any_name, qualified_name, unquote, }; use crate::models::{ CheckConstraint, Column, ColumnGenerated, ConstraintColumns, ForeignKey, - ForeignKeyReference, Index, IndexColumn, Table, + ForeignKeyReference, Index, IndexColumn, LikeTable, Table, TablePartition, + TablePartitionBehavior, TablePartitionColumn, }; use crate::utils::quote_ident; -/// CREATE TABLE → Table (columns + inline constraints) +/// CREATE TABLE → Table (columns + inline constraints), or a +/// `PARTITION OF` child, returned as [`Statement::CreateTablePartition`] +/// so the pull assembly can fold it into its parent's `partitions` pub(crate) fn create_table( node: &Node, src: &str, ) -> Result { - let name = node - .find("qualified_name") - .ok_or_else(|| String::from("CREATE TABLE without a name"))?; - let name = qualified_name(&name, src)?; + // two variants carry a second qualified_name: `PARTITION OF parent` + // (handled below) and inline `REFERENCES` (nested inside a column + // constraint, so it never surfaces as a *direct* child); find_all + // here only needs the table's own name plus, for PARTITION OF, the + // parent's + let names = node.find_all("qualified_name"); + let name = match names.first() { + Some(n) => qualified_name(n, src)?, + None => return Err(String::from("CREATE TABLE without a name")), + }; + // `PARTITION OF parent FOR VALUES ...` — a partition child, not a + // standalone table. The kw_partition + kw_of combination is unique + // to this form (PARTITION BY uses kw_partition + kw_by instead) + if node.has("kw_partition") && node.has("kw_of") { + let parent = match names.get(1) { + Some(n) => qualified_name(n, src)?, + None => { + return Err(String::from( + "PARTITION OF without a parent table", + )); + } + }; + let bound = + node.child_of_kind("PartitionBoundSpec").ok_or_else(|| { + String::from("PARTITION OF without a FOR VALUES clause") + })?; + let bound = partition_bound(&bound, src); + return Ok(Statement::CreateTablePartition { + parent, + partition: TablePartition { + name: name.name, + schema: name.schema.unwrap_or_default(), + default: bound.default, + for_values_in: bound.for_values_in, + for_values_from: bound.for_values_from, + for_values_to: bound.for_values_to, + for_values_with: bound.for_values_with, + comment: None, + }, + }); + } let mut table = Table { name: name.name, schema: name.schema.unwrap_or_default(), owner: String::new(), sql: None, unlogged: node.has("kw_unlogged").then_some(true), - from_type: None, + // `CREATE TABLE name OF typename` — direct child only, since + // `any_name` also appears nested inside Typename for other + // CreateStmt forms + from_type: node + .has("kw_of") + .then(|| { + node.child_of_kind("any_name").map(|n| { + let of_type = any_name(&n, src); + match of_type.schema { + Some(schema) => format!( + "{}.{}", + quote_ident(&schema), + quote_ident(&of_type.name) + ), + None => quote_ident(&of_type.name), + } + }) + }) + .flatten(), parents: None, like_table: None, columns: None, @@ -36,11 +96,19 @@ pub(crate) fn create_table( unique_constraints: None, foreign_keys: None, triggers: None, - partition: None, + partition: table_partition_behavior(node, src), partitions: None, - access_method: None, - storage_parameters: None, - tablespace: None, + access_method: node + .child_of_kind("table_access_method_clause") + .and_then(|n| n.child_of_kind("name")) + .map(|n| unquote(n.text(src))), + storage_parameters: node + .child_of_kind("OptWith") + .and_then(|n| reloptions(&n, src)), + tablespace: node + .child_of_kind("OptTableSpace") + .and_then(|n| n.find("name")) + .map(|n| unquote(n.text(src))), index_tablespace: None, server: None, options: None, @@ -53,8 +121,21 @@ pub(crate) fn create_table( } else if let Some(constraint) = element.child_of_kind("TableConstraint") { + // USING INDEX TABLESPACE, on PRIMARY KEY/UNIQUE/EXCLUDE — + // the model keeps a single table-level index_tablespace, so + // the first constraint that carries one wins + if table.index_tablespace.is_none() { + table.index_tablespace = constraint + .find("OptConsTableSpace") + .and_then(|n| n.find("name")) + .map(|n| unquote(n.text(src))); + } let (name, parsed) = table_constraint(&constraint, src)?; apply_constraint(&mut table, name, parsed); + } else if let Some(like_clause) = + element.child_of_kind("TableLikeClause") + { + table.like_table = Some(like_table(&like_clause, src)); } } if !columns.is_empty() { @@ -188,7 +269,9 @@ pub(crate) fn create_index( .child_of_kind("where_clause") .and_then(|n| n.find("a_expr")) .map(|n| n.text(src).to_string()), - storage_parameters: None, + storage_parameters: node + .child_of_kind("opt_reloptions") + .and_then(|n| reloptions(&n, src)), tablespace: node .child_of_kind("OptTableSpace") .and_then(|n| n.find("name")) @@ -349,6 +432,19 @@ fn foreign_key( .collect() }) .unwrap_or_default(); + let spec = elem.child_of_kind("ConstraintAttributeSpec"); + let deferrable = spec.and_then(|s| { + s.find_all("ConstraintAttributeElem") + .iter() + .find(|e| e.has("kw_deferrable")) + .map(|e| !e.has("kw_not")) + }); + let initially_deferred = spec.and_then(|s| { + s.find_all("ConstraintAttributeElem") + .iter() + .find(|e| e.has("kw_initially")) + .map(|e| e.has("kw_deferred")) + }); let mut on_delete = None; let mut on_update = None; if let Some(actions) = elem.child_of_kind("key_actions") { @@ -382,8 +478,8 @@ fn foreign_key( }), on_delete, on_update, - deferrable: None, - initially_deferred: None, + deferrable, + initially_deferred, }) } @@ -417,6 +513,209 @@ pub(crate) fn apply_constraint( } } +/// `PARTITION BY ()`, if present +fn table_partition_behavior( + node: &Node, + src: &str, +) -> Option { + let spec = node.find("PartitionSpec")?; + let partition_type = spec + .child_of_kind("ColId") + .map(|n| n.text(src).to_uppercase()) + .unwrap_or_default(); + let columns: Vec = spec + .find_all("part_elem") + .iter() + .map(|elem| partition_column(elem, src)) + .collect(); + (!columns.is_empty()).then_some(TablePartitionBehavior { + partition_type, + columns, + }) +} + +fn partition_column(elem: &Node, src: &str) -> TablePartitionColumn { + let name = elem.child_of_kind("ColId").map(|n| unquote(n.text(src))); + let expression = if name.is_none() { + elem.child_of_kind("func_expr_windowless") + .or_else(|| elem.child_of_kind("a_expr")) + .map(|n| n.text(src).to_string()) + } else { + None + }; + let collation = elem + .find("opt_collate") + .and_then(|n| n.find("any_name")) + .map(|n| n.text(src).to_string()); + let opclass = elem + .find("opt_qualified_name") + .map(|n| n.text(src).to_string()); + match (&name, &collation, &opclass) { + (Some(name), None, None) => TablePartitionColumn::Name(name.clone()), + _ => TablePartitionColumn::Detailed { + name, + expression, + collation, + opclass, + }, + } +} + +/// The bound of a `PARTITION OF ...` child, one field of which is set +/// depending on the `PartitionBoundSpec` alternative matched +#[derive(Default)] +struct PartitionBound { + default: Option, + for_values_in: Option>, + for_values_from: Option, + for_values_to: Option, + for_values_with: Option, +} + +fn partition_bound(spec: &Node, src: &str) -> PartitionBound { + if spec.has("kw_default") { + return PartitionBound { + default: Some(true), + ..Default::default() + }; + } + if spec.has("kw_with") { + return PartitionBound { + for_values_with: spec + .child_of_kind("hash_partbound") + .map(|n| n.text(src).to_string()), + ..Default::default() + }; + } + if spec.has("kw_from") { + let lists = spec.find_all("expr_list"); + return PartitionBound { + for_values_from: lists.first().map(|n| partition_value(n, src)), + for_values_to: lists.get(1).map(|n| partition_value(n, src)), + ..Default::default() + }; + } + if spec.has("kw_in") { + let values = spec + .child_of_kind("expr_list") + .map(|n| partition_values(&n, src)) + .unwrap_or_default(); + return PartitionBound { + for_values_in: (!values.is_empty()).then_some(values), + ..Default::default() + }; + } + PartitionBound::default() +} + +fn partition_value(list: &Node, src: &str) -> Value { + let exprs = list.find_all("a_expr"); + match exprs.as_slice() { + [one] => single_expr_value(one, src), + _ => Value::Array( + exprs.iter().map(|e| single_expr_value(e, src)).collect(), + ), + } +} + +fn partition_values(list: &Node, src: &str) -> Vec { + list.find_all("a_expr") + .iter() + .map(|e| single_expr_value(e, src)) + .collect() +} + +fn single_expr_value(node: &Node, src: &str) -> Value { + if let Some(s) = node.find("Sconst") { + Value::String(string_value(&s, src)) + } else if let Ok(n) = node.text(src).parse::() { + Value::Number(n.into()) + } else { + Value::String(node.text(src).to_string()) + } +} + +/// `LIKE source_table [{INCLUDING|EXCLUDING} option ...]` +fn like_table(clause: &Node, src: &str) -> LikeTable { + let name = clause + .child_of_kind("qualified_name") + .and_then(|n| qualified_name(&n, src).ok()) + .map(|q| match q.schema { + Some(schema) => { + format!("{}.{}", quote_ident(&schema), quote_ident(&q.name)) + } + None => quote_ident(&q.name), + }) + .unwrap_or_default(); + let mut like = LikeTable { + name, + include_comments: None, + include_constraints: None, + include_defaults: None, + include_generated: None, + include_identity: None, + include_indexes: None, + include_statistics: None, + include_storage: None, + include_all: None, + }; + if let Some(list) = clause.child_of_kind("TableLikeOptionList") { + for (including, option) in like_options(&list) { + match option { + "COMMENTS" => like.include_comments = Some(including), + "CONSTRAINTS" => like.include_constraints = Some(including), + "DEFAULTS" => like.include_defaults = Some(including), + "GENERATED" => like.include_generated = Some(including), + "IDENTITY" => like.include_identity = Some(including), + "INDEXES" => like.include_indexes = Some(including), + "STATISTICS" => like.include_statistics = Some(including), + "STORAGE" => like.include_storage = Some(including), + "ALL" => like.include_all = Some(including), + _ => {} + } + } + } + like +} + +/// Walk the left-recursive `TableLikeOptionList` chain into +/// `(including, OPTION_NAME)` pairs +fn like_options(node: &Node) -> Vec<(bool, &'static str)> { + let mut results = Vec::new(); + if let Some(inner) = node.child_of_kind("TableLikeOptionList") { + results.extend(like_options(&inner)); + } + if let Some(option) = node.child_of_kind("TableLikeOption") { + let including = node.child_of_kind("kw_including").is_some(); + results.push((including, like_option_name(&option))); + } + results +} + +fn like_option_name(node: &Node) -> &'static str { + if node.has("kw_comments") { + "COMMENTS" + } else if node.has("kw_constraints") { + "CONSTRAINTS" + } else if node.has("kw_defaults") { + "DEFAULTS" + } else if node.has("kw_generated") { + "GENERATED" + } else if node.has("kw_identity") { + "IDENTITY" + } else if node.has("kw_indexes") { + "INDEXES" + } else if node.has("kw_statistics") { + "STATISTICS" + } else if node.has("kw_storage") { + "STORAGE" + } else if node.has("kw_all") { + "ALL" + } else { + "" + } +} + #[cfg(test)] mod tests { use serde_json::json; @@ -660,4 +959,152 @@ mod tests { let statement = parse_one("VACUUM ANALYZE test.users;"); assert!(matches!(statement, Statement::Unsupported(_))); } + + #[test] + fn parses_partition_by() { + let statement = parse_one( + "CREATE TABLE test.events (id bigint, ts timestamp) \ + PARTITION BY RANGE (ts);", + ); + let Statement::CreateTable(table) = statement else { + panic!("expected CreateTable") + }; + let partition = table.partition.unwrap(); + assert_eq!(partition.partition_type, "RANGE"); + assert_eq!( + partition.columns, + vec![TablePartitionColumn::Name("ts".into())] + ); + } + + #[test] + fn parses_partition_of() { + let statement = parse_one( + "CREATE TABLE test.events_2024 PARTITION OF test.events \ + FOR VALUES FROM ('2024-01-01') TO ('2025-01-01');", + ); + let Statement::CreateTablePartition { parent, partition } = statement + else { + panic!("expected CreateTablePartition") + }; + assert_eq!(parent.to_string(), "test.events"); + assert_eq!(partition.schema, "test"); + assert_eq!(partition.name, "events_2024"); + assert_eq!(partition.for_values_from, Some(json!("2024-01-01"))); + assert_eq!(partition.for_values_to, Some(json!("2025-01-01"))); + } + + #[test] + fn parses_partition_of_multicolumn_bounds() { + let statement = parse_one( + "CREATE TABLE test.events_2024 PARTITION OF test.events \ + FOR VALUES FROM (2020, 1) TO (2021, 1);", + ); + let Statement::CreateTablePartition { partition, .. } = statement + else { + panic!("expected CreateTablePartition") + }; + assert_eq!(partition.for_values_from, Some(json!([2020, 1]))); + assert_eq!(partition.for_values_to, Some(json!([2021, 1]))); + } + + #[test] + fn parses_partition_of_default() { + let statement = parse_one( + "CREATE TABLE test.events_default PARTITION OF test.events \ + DEFAULT;", + ); + let Statement::CreateTablePartition { partition, .. } = statement + else { + panic!("expected CreateTablePartition") + }; + assert_eq!(partition.default, Some(true)); + } + + #[test] + fn parses_create_table_of_type() { + let statement = + parse_one("CREATE TABLE test.person OF test.person_type;"); + let Statement::CreateTable(table) = statement else { + panic!("expected CreateTable") + }; + assert_eq!(table.from_type, Some("test.person_type".into())); + } + + #[test] + fn parses_create_table_like() { + let statement = parse_one( + "CREATE TABLE test.copy (\n\ + LIKE test.original INCLUDING DEFAULTS INCLUDING INDEXES\n\ + );", + ); + let Statement::CreateTable(table) = statement else { + panic!("expected CreateTable") + }; + let like_table = table.like_table.unwrap(); + assert_eq!(like_table.name, "test.original"); + assert_eq!(like_table.include_defaults, Some(true)); + assert_eq!(like_table.include_indexes, Some(true)); + } + + #[test] + fn parses_table_storage_options() { + let statement = parse_one( + "CREATE TABLE test.t (id int) USING heap \ + WITH (fillfactor=70) TABLESPACE fastdisk;", + ); + let Statement::CreateTable(table) = statement else { + panic!("expected CreateTable") + }; + assert_eq!(table.access_method, Some("heap".into())); + assert_eq!( + table.storage_parameters.unwrap().get("fillfactor"), + Some(&json!("70")) + ); + assert_eq!(table.tablespace, Some("fastdisk".into())); + } + + #[test] + fn parses_primary_key_index_tablespace() { + let statement = parse_one( + "CREATE TABLE test.t (\n\ + id int,\n\ + CONSTRAINT t_pkey PRIMARY KEY (id) USING INDEX TABLESPACE fast\n\ + );", + ); + let Statement::CreateTable(table) = statement else { + panic!("expected CreateTable") + }; + assert_eq!(table.index_tablespace, Some("fast".into())); + } + + #[test] + fn parses_index_storage_parameters() { + let statement = + parse_one("CREATE INDEX i ON t (c) WITH (fillfactor=80);"); + let Statement::CreateIndex { index, .. } = statement else { + panic!("expected CreateIndex") + }; + assert_eq!( + index.storage_parameters.unwrap().get("fillfactor"), + Some(&json!("80")) + ); + } + + #[test] + fn parses_foreign_key_deferrable() { + let statement = parse_one( + "ALTER TABLE ONLY test.addresses\n \ + ADD CONSTRAINT addresses_user_id_fkey FOREIGN KEY (user_id) \ + REFERENCES test.users(id) DEFERRABLE INITIALLY DEFERRED;", + ); + let Statement::AddConstraint { constraint, .. } = statement else { + panic!("expected AddConstraint") + }; + let TableConstraint::ForeignKey(fk) = constraint else { + panic!("expected ForeignKey") + }; + assert_eq!(fk.deferrable, Some(true)); + assert_eq!(fk.initially_deferred, Some(true)); + } } diff --git a/src/ddl/view.rs b/src/ddl/view.rs index f67fe1c..2cc25ed 100644 --- a/src/ddl/view.rs +++ b/src/ddl/view.rs @@ -2,6 +2,7 @@ use tree_sitter::Node; +use crate::ddl::object::reloptions; use crate::ddl::{NodeExt, Statement, qualified_name, unquote}; use crate::models::{MaterializedView, View, ViewColumn}; @@ -33,7 +34,14 @@ pub(crate) fn create_view( } .to_string() }), - security_barrier: None, + security_barrier: node + .child_of_kind("opt_reloptions") + .and_then(|n| reloptions(&n, src)) + .and_then(|m| m.get("security_barrier").cloned()) + .and_then(|v| match v { + serde_json::Value::String(s) => pg_bool(&s), + _ => None, + }), query, comment: None, })) @@ -65,7 +73,9 @@ pub(crate) fn create_materialized_view( .find("table_access_method_clause") .and_then(|n| n.find("name")) .map(|n| unquote(n.text(src))), - storage_parameters: None, + storage_parameters: node + .find("opt_reloptions") + .and_then(|n| reloptions(&n, src)), tablespace: node .find("OptTableSpace") .and_then(|n| n.find("name")) @@ -90,6 +100,18 @@ fn view_columns(node: &Node, src: &str) -> Option> { (!columns.is_empty()).then_some(columns) } +/// Parse a PostgreSQL boolean reloption value. PostgreSQL accepts +/// `true/t/on/yes/y/1` and `false/f/off/no/n/0` (case-insensitively) +/// for boolean reloptions; a bare option key round-trips through +/// `reloptions` as `"true"`. +fn pg_bool(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "true" | "t" | "on" | "yes" | "y" | "1" => Some(true), + "false" | "f" | "off" | "no" | "n" | "0" => Some(false), + _ => None, + } +} + #[cfg(test)] mod tests { use super::*; @@ -145,4 +167,59 @@ mod tests { assert_eq!(view.name, "mv"); assert_eq!(view.query, Some("SELECT id FROM test.users".into())); } + + #[test] + fn parses_view_security_barrier() { + let Statement::CreateView(view) = parse_one( + "CREATE VIEW test.v WITH (security_barrier=true) AS \ + SELECT 1;", + ) else { + panic!("expected CreateView") + }; + assert_eq!(view.security_barrier, Some(true)); + } + + #[test] + fn parses_view_bare_security_barrier() { + let Statement::CreateView(view) = parse_one( + "CREATE VIEW test.v WITH (security_barrier) AS SELECT 1;", + ) else { + panic!("expected CreateView") + }; + assert_eq!(view.security_barrier, Some(true)); + } + + #[test] + fn parses_view_security_barrier_boolean_forms() { + for (sql, expected) in [ + ("WITH (security_barrier=on)", Some(true)), + ("WITH (security_barrier=off)", Some(false)), + ("WITH (security_barrier='no')", Some(false)), + ("WITH (security_barrier='t')", Some(true)), + ("WITH (security_barrier='f')", Some(false)), + ("WITH (security_barrier='y')", Some(true)), + ("WITH (security_barrier='n')", Some(false)), + ] { + let stmt = + parse_one(&format!("CREATE VIEW test.v {sql} AS SELECT 1;")); + let Statement::CreateView(view) = stmt else { + panic!("expected CreateView") + }; + assert_eq!(view.security_barrier, expected, "for {sql}"); + } + } + + #[test] + fn parses_materialized_view_storage_parameters() { + let Statement::CreateMaterializedView(view) = parse_one( + "CREATE MATERIALIZED VIEW test.mv WITH (fillfactor=90) AS \ + SELECT id FROM test.users;", + ) else { + panic!("expected CreateMaterializedView") + }; + assert_eq!( + view.storage_parameters.unwrap().get("fillfactor"), + Some(&serde_json::Value::String("90".into())) + ); + } } diff --git a/src/deploy/alter.rs b/src/deploy/alter.rs index 5e14fea..16d37be 100644 --- a/src/deploy/alter.rs +++ b/src/deploy/alter.rs @@ -14,8 +14,8 @@ use crate::build; use crate::deploy::diff::canonical_type; use crate::models::{ CheckConstraint, Column, Definition, Domain, Extension, - ForeignDataWrapper, ForeignKey, Index, Schema, Sequence, Server, Table, - Trigger, Type, UserMapping, View, ViewColumn, + ForeignDataWrapper, ForeignKey, Function, Index, Schema, Sequence, Server, + Table, Trigger, Type, UserMapping, View, ViewColumn, }; use crate::utils::{postgres_value, quote_ident, user_mapping_subject}; @@ -70,11 +70,14 @@ pub(crate) fn resolve(repo: &Definition, database: &Definition) -> Resolution { extension(repo, db) } (Definition::Schema(repo), Definition::Schema(db)) => schema(repo, db), - // CREATE OR REPLACE handles function bodies and view queries - // in place; a function whose return type changed cannot be - // replaced and must be dropped first + // CREATE OR REPLACE handles function bodies and view queries in + // place; a function whose return type or output-parameter + // signature changed cannot be replaced (Postgres rejects an + // OR REPLACE that alters the output) and must be dropped first (Definition::Function(repo), Definition::Function(db)) => { - if repo.returns == db.returns { + if returns_equal(repo, db) + && out_parameters(repo) == out_parameters(db) + { // function names carry their full identity signature, // so COMMENT ON FUNCTION takes the name verbatim let target = @@ -108,6 +111,35 @@ fn qualified(schema: &str, name: &str) -> String { format!("{}.{}", quote_ident(schema), quote_ident(name)) } +/// True when two functions' return types are the same modulo type +/// aliasing (`int4` vs `integer`) +fn returns_equal(repo: &Function, db: &Function) -> bool { + match (&repo.returns, &db.returns) { + (Some(r), Some(d)) => canonical_type(r) == canonical_type(d), + (r, d) => r == d, + } +} + +/// The `OUT`/`TABLE`-mode parameters that make up a function's output +/// signature, with types canonicalized so an alias does not spuriously +/// diff. `CREATE OR REPLACE FUNCTION` cannot change this signature, so +/// callers must fall back to a drop+recreate when it differs. +fn out_parameters(function: &Function) -> Vec<(String, String, String)> { + function + .parameters + .iter() + .flatten() + .filter(|p| p.mode == "OUT" || p.mode == "TABLE") + .map(|p| { + ( + p.mode.clone(), + p.name.clone().unwrap_or_default(), + canonical_type(&p.data_type), + ) + }) + .collect() +} + /// CREATE OR REPLACE VIEW only succeeds when the new query's output /// columns stay compatible with the existing view: same names in the /// same order, with new columns added only at the end. A rename, @@ -555,8 +587,12 @@ fn sequence(repo: &Sequence, db: &Sequence) -> Resolution { /// place. A base-type, collation, or constraint change rebuilds (the /// domain's constraints are not all individually named). fn domain(repo: &Domain, db: &Domain) -> Resolution { + let data_type_changed = match (&repo.data_type, &db.data_type) { + (Some(r), Some(d)) => canonical_type(r) != canonical_type(d), + (r, d) => r != d, + }; if repo.sql != db.sql - || repo.data_type != db.data_type + || data_type_changed || repo.collation != db.collation || repo.check_constraints != db.check_constraints { @@ -600,8 +636,10 @@ fn enum_type(repo: &Type, db: &Type) -> Resolution { let mut alters: Vec = repo_values[db_values.len()..] .iter() .map(|value| { - let escaped = value.replace('\'', "''"); - Alter::new(format!("ALTER TYPE {name} ADD VALUE '{escaped}';\n")) + Alter::new(format!( + "ALTER TYPE {name} ADD VALUE {};\n", + string_literal(value) + )) }) .collect(); if repo.comment != db.comment { @@ -622,7 +660,8 @@ fn extension(repo: &Extension, db: &Extension) -> Resolution { && let Some(version) = &repo.version { alters.push(Alter::new(format!( - "ALTER EXTENSION {name} UPDATE TO '{version}';\n" + "ALTER EXTENSION {name} UPDATE TO {};\n", + string_literal(version) ))); } if repo.schema != db.schema @@ -666,7 +705,8 @@ fn schema(repo: &Schema, db: &Schema) -> Resolution { /// use the FOREIGN TABLE object type, matching the build. fn foreign_table(repo: &Table, db: &Table) -> Resolution { if repo.server != db.server - || repo.columns != db.columns + || canonicalize_columns(&repo.columns) + != canonicalize_columns(&db.columns) || repo.sql != db.sql || repo.check_constraints != db.check_constraints { @@ -689,6 +729,20 @@ fn foreign_table(repo: &Table, db: &Table) -> Resolution { Resolution::Statements(alters) } +/// A copy of `columns` with each data type canonicalized, so an alias +/// (`int4` vs `integer`) does not force an unnecessary rebuild +fn canonicalize_columns(columns: &Option>) -> Vec { + columns + .iter() + .flatten() + .cloned() + .map(|mut column| { + column.data_type = canonical_type(&column.data_type); + column + }) + .collect() +} + /// Foreign-data-wrapper reconciliation: handler, validator, OPTIONS, /// and comment are all alterable in place. fn fdw(repo: &ForeignDataWrapper, db: &ForeignDataWrapper) -> Resolution { @@ -1272,6 +1326,72 @@ mod tests { )); } + #[test] + fn function_returns_alias_uses_or_replace() { + let f = |returns: &str| -> Definition { + Definition::Function( + serde_json::from_value(serde_json::json!({ + "name": "f", "schema": "test", "owner": "postgres", + "returns": returns, "language": "sql", + "definition": "SELECT 1", + })) + .unwrap(), + ) + }; + // a repo `returns: int4` against the server's `integer` must + // not force a drop+recreate + assert!(matches!( + resolve(&f("int4"), &f("integer")), + Resolution::OrReplace { .. } + )); + } + + #[test] + fn function_out_parameter_change_replaces() { + let f = |data_type: &str| -> Definition { + Definition::Function( + serde_json::from_value(serde_json::json!({ + "name": "f", "schema": "test", "owner": "postgres", + "returns": "record", "language": "sql", + "definition": "SELECT 1", + "parameters": [ + {"mode": "OUT", "name": "x", "data_type": data_type}, + ], + })) + .unwrap(), + ) + }; + // CREATE OR REPLACE FUNCTION cannot change the output + // signature, so a changed OUT parameter must fall back to a + // gated drop+recreate rather than OrReplace + assert!(matches!( + resolve(&f("bigint"), &f("integer")), + Resolution::Replace + )); + } + + #[test] + fn function_out_parameter_alias_uses_or_replace() { + let f = |data_type: &str| -> Definition { + Definition::Function( + serde_json::from_value(serde_json::json!({ + "name": "f", "schema": "test", "owner": "postgres", + "returns": "record", "language": "sql", + "definition": "SELECT 1", + "parameters": [ + {"mode": "OUT", "name": "x", "data_type": data_type}, + ], + })) + .unwrap(), + ) + }; + // an aliased OUT parameter type must not force a rebuild + assert!(matches!( + resolve(&f("int4"), &f("integer")), + Resolution::OrReplace { .. } + )); + } + #[test] fn view_change_uses_or_replace() { let v = |query: &str| -> Definition { @@ -1498,6 +1618,21 @@ mod tests { assert!(matches!(domain(&retyped, &db), Resolution::Replace)); } + #[test] + fn domain_type_alias_is_not_replace() { + let db: Domain = serde_json::from_value(serde_json::json!({ + "name": "d", "schema": "test", "owner": "postgres", + "data_type": "integer", + })) + .unwrap(); + let mut repo = db.clone(); + repo.data_type = Some("int4".into()); + assert!(matches!( + domain(&repo, &db), + Resolution::Statements(ref alters) if alters.is_empty() + )); + } + #[test] fn schema_comment_changes_in_place() { let db: Schema = serde_json::from_value(serde_json::json!({ @@ -1717,6 +1852,26 @@ mod tests { assert!(alters.iter().all(|a| !a.destructive)); } + #[test] + fn foreign_table_column_type_alias_is_not_replace() { + let repo = parse_table(foreign_table_value( + serde_json::json!({"table_name": "t"}), + None, + )); + let mut db = repo.clone(); + db.columns = Some(vec![Column { + name: "id".into(), + data_type: "int4".into(), + nullable: None, + default: None, + collation: None, + check_constraint: None, + generated: None, + comment: None, + }]); + assert!(matches!(table(&repo, &db), Resolution::Statements(_))); + } + #[test] fn foreign_table_server_change_replaces() { let repo = parse_table(foreign_table_value( diff --git a/src/deploy/diff.rs b/src/deploy/diff.rs index 9b93f3c..b488c2f 100644 --- a/src/deploy/diff.rs +++ b/src/deploy/diff.rs @@ -282,7 +282,7 @@ fn normalize(value: &mut Value) { // cannot apply pg_restore-style ownership anyway) map.remove("owner"); for (key, child) in map.iter_mut() { - if key == "data_type" + if (key == "data_type" || key == "returns") && let Some(data_type) = child.as_str() { *child = Value::String(canonical_type(data_type)); @@ -410,6 +410,26 @@ mod tests { ); } + #[test] + fn function_returns_alias_is_not_a_change() { + let f = |returns: &str| -> Definition { + Definition::Function( + serde_json::from_value(serde_json::json!({ + "name": "f", + "schema": "test", + "owner": "postgres", + "returns": returns, + "language": "sql", + "definition": "SELECT 1", + })) + .unwrap(), + ) + }; + // a repo `returns: int4` must not diff against the server's + // `integer` on every deploy + assert_eq!(normalized(&f("int4")), normalized(&f("integer"))); + } + #[test] fn function_key_canonicalizes_parameter_types() { let repo: crate::models::Function = diff --git a/src/pull/mod.rs b/src/pull/mod.rs index 1a44339..ebbf63a 100644 --- a/src/pull/mod.rs +++ b/src/pull/mod.rs @@ -392,6 +392,9 @@ pub struct Assembly { /// index entry was seen (pg_dump sorts INDEX before MATERIALIZED /// VIEW), replayed after the entry loop completes deferred_indexes: Vec<(QualifiedName, models::Index)>, + /// PARTITION OF children whose parent table had not yet been + /// ingested, replayed after the entry loop completes + deferred_partitions: Vec<(QualifiedName, models::TablePartition)>, } impl Assembly { @@ -528,6 +531,7 @@ impl Assembly { } } self.apply_deferred_indexes(); + self.apply_deferred_partitions(); task.finish(); Ok(()) } @@ -546,6 +550,22 @@ impl Assembly { } } + /// Attach PARTITION OF children whose parent table was not yet + /// ingested when the child entry was seen; warn for any that + /// remain unresolved + fn apply_deferred_partitions(&mut self) { + for (parent, partition) in + std::mem::take(&mut self.deferred_partitions) + { + match self.find_table(&parent) { + Some(table) => { + table.partitions.get_or_insert_default().push(partition); + } + None => log::warn!("Partition of unknown table {parent}"), + } + } + } + /// Parse a `pg_dumpall --roles-only` SQL dump, skipping comments, /// SET statements, and psql meta-commands (PG17 wraps the output /// in `\restrict` / `\unrestrict`) @@ -597,6 +617,19 @@ impl Assembly { table.owner = owner; self.tables.push(*table); } + Statement::CreateTablePartition { parent, partition } => { + match self.find_table(&parent) { + Some(table) => { + table + .partitions + .get_or_insert_default() + .push(partition); + } + None => { + self.deferred_partitions.push((parent, partition)); + } + } + } Statement::CreateSequence(mut sequence) => { sequence.owner = owner; self.sequences.push(sequence); @@ -1497,6 +1530,43 @@ mod tests { assert_eq!(fks[0].on_delete.as_deref(), Some("CASCADE")); } + #[test] + fn merges_partition_children() { + // the child entry (events_2024) sorts before its parent + // (events) in this dump, exercising the deferred-partition + // retry path the same way deferred indexes are exercised + let mut dump = libpgdump::new("fixtures", "UTF8", "18.0").unwrap(); + add(&mut dump, OT::Schema, "", "test", "CREATE SCHEMA test;"); + add( + &mut dump, + OT::Table, + "test", + "events_2024", + "CREATE TABLE test.events_2024 PARTITION OF test.events \ + FOR VALUES FROM ('2024-01-01') TO ('2025-01-01');", + ); + add( + &mut dump, + OT::Table, + "test", + "events", + "CREATE TABLE test.events (id bigint, ts timestamp) \ + PARTITION BY RANGE (ts);", + ); + let mut assembly = Assembly::default(); + assembly.ingest(&dump).unwrap(); + assert_eq!(assembly.tables.len(), 1); + let events = &assembly.tables[0]; + assert_eq!(events.name, "events"); + let partitions = events.partitions.as_ref().expect("partitions"); + assert_eq!(partitions.len(), 1); + assert_eq!(partitions[0].name, "events_2024"); + assert_eq!( + partitions[0].for_values_from, + Some(serde_json::json!("2024-01-01")) + ); + } + #[test] fn merges_sequence_owned_by() { let assembly = assembled(); diff --git a/src/pull/update.rs b/src/pull/update.rs index 0f3f0a0..83bcdf4 100644 --- a/src/pull/update.rs +++ b/src/pull/update.rs @@ -56,7 +56,7 @@ pub fn merge( .map_err(|e| format!("failed to write {}: {e}", path.display()))?; written += 1; } - let stale = stale_files(root, files, &ignore)?; + let stale = stale_files(root, files, &ignore, args)?; if args.prune { for relative in &stale { log::info!("Removing {}", relative.display()); @@ -86,6 +86,99 @@ pub fn merge( /// staleness like the directories above const REMAINING_FILE: &str = "remaining.yaml"; +/// Whether this run extracted cluster roles/users via `pg_dumpall`. +/// Mirrors the gate in `pull::pull` (`roles = None` when `--no-roles` +/// or `--dump` is given): if it did not run, the render produced no +/// role/user files, and staleness must not be judged against `roles/` +/// and `users/` or every pre-existing file there would be pruned +fn roles_extracted(args: &cli::Pull) -> bool { + !args.no_roles && args.dump.is_none() +} + +/// The schema an on-disk managed-directory file belongs to, so it can +/// be checked against an active `--exclude-schema` filter. `dir` is +/// the top-level managed directory (e.g. `"tables"`) and `relative` +/// the file's path under `root` +fn schema_of(dir: &str, relative: &Path) -> Option { + match dir { + "schemata" | "types" => relative + .file_stem() + .map(|s| s.to_string_lossy().into_owned()), + "domains" | "functions" | "materialized_views" | "sequences" + | "tables" | "views" => relative + .strip_prefix(dir) + .ok() + .and_then(|rest| rest.iter().next()) + .map(|s| s.to_string_lossy().into_owned()), + _ => None, + } +} + +/// The unqualified object name for a file under a directory +/// `--exclude-table` applies to (tables, views, materialized views, +/// sequences), for checking against an active `--exclude-table` filter +fn table_name_of(dir: &str, relative: &Path) -> Option { + match dir { + "materialized_views" | "sequences" | "tables" | "views" => relative + .file_stem() + .map(|s| s.to_string_lossy().into_owned()), + _ => None, + } +} + +/// Whether `text` matches a `pg_dump`-style pattern (`*` and `?` +/// wildcards; `--exclude-schema`/`--exclude-table` patterns) +fn glob_match(pattern: &str, text: &str) -> bool { + let p: Vec = pattern.chars().collect(); + let t: Vec = text.chars().collect(); + let (mut pi, mut ti) = (0, 0); + let (mut star, mut resume) = (None, 0); + while ti < t.len() { + if pi < p.len() && (p[pi] == '?' || p[pi] == t[ti]) { + pi += 1; + ti += 1; + } else if pi < p.len() && p[pi] == '*' { + star = Some(pi); + resume = ti; + pi += 1; + } else if let Some(si) = star { + pi = si + 1; + resume += 1; + ti = resume; + } else { + return false; + } + } + while p.get(pi) == Some(&'*') { + pi += 1; + } + pi == p.len() +} + +/// Whether a file the render did not produce is still covered by an +/// active `--exclude-schema`/`--exclude-table` filter — such objects +/// were deliberately left out of this run's dump, so their files are +/// not stale +fn filtered_out(dir: &str, relative: &Path, args: &cli::Pull) -> bool { + if let Some(schema) = schema_of(dir, relative) + && args.exclude_schema.iter().any(|p| glob_match(p, &schema)) + { + return true; + } + if let Some(name) = table_name_of(dir, relative) { + let schema = schema_of(dir, relative).unwrap_or_default(); + let qualified = format!("{schema}.{name}"); + if args + .exclude_table + .iter() + .any(|p| glob_match(p, &qualified) || glob_match(p, &name)) + { + return true; + } + } + false +} + /// YAML files on disk under the managed directories (plus the /// root-level `remaining.yaml`) that the render did not produce and /// the ignore file does not cover @@ -93,6 +186,7 @@ fn stale_files( root: &Path, files: &BTreeMap, ignore: &BTreeSet, + args: &cli::Pull, ) -> Result, String> { let mut stale = Vec::new(); let remaining = PathBuf::from(REMAINING_FILE); @@ -102,15 +196,19 @@ fn stale_files( { stale.push(remaining); } + let roles_extracted = roles_extracted(args); for dir in MANAGED_DIRS { + if (*dir == "roles" || *dir == "users") && !roles_extracted { + continue; + } let top = root.join(dir); if !top.is_dir() { continue; } let mut pending = vec![top]; - while let Some(dir) = pending.pop() { - for entry in std::fs::read_dir(&dir).map_err(|e| { - format!("failed to read {}: {e}", dir.display()) + while let Some(dir_path) = pending.pop() { + for entry in std::fs::read_dir(&dir_path).map_err(|e| { + format!("failed to read {}: {e}", dir_path.display()) })? { let entry = entry.map_err(|e| format!("failed to read entry: {e}"))?; @@ -134,6 +232,7 @@ fn stale_files( .to_path_buf(); if files.contains_key(&relative) || ignore.contains(&relative.to_string_lossy().to_string()) + || filtered_out(dir, &relative, args) { continue; } @@ -189,3 +288,133 @@ fn remove_emptied_directories(root: &Path) -> Result<(), String> { } Ok(()) } + +#[cfg(test)] +mod tests { + use std::fs; + + use clap::Parser; + + use super::*; + + /// Build `cli::Pull` args for `pull --update --prune ` + /// the way the binary would parse them + fn pull_args(extra: &[&str], dest: &Path) -> cli::Pull { + let mut argv = vec!["pglifecycle", "pull", "--update", "--prune"]; + argv.extend_from_slice(extra); + let dest = dest.to_str().unwrap(); + argv.push(dest); + match cli::Cli::try_parse_from(argv).unwrap().action { + cli::Action::Pull(args) => args, + _ => unreachable!(), + } + } + + fn touch(root: &Path, relative: &str) { + let path = root.join(relative); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(path, "").unwrap(); + } + + /// `--dump` skips pg_dumpall (mirrors the gate in `pull::pull`), so + /// a render from a dump produces no role/user files; staleness must + /// not treat pre-existing `roles/`/`users/` files as stale, or + /// `--prune` would wipe the entire cluster role inventory + #[test] + fn dump_source_does_not_stale_roles_or_users() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + touch(root, "roles/app_group.yaml"); + touch(root, "users/app_login.yaml"); + let dump = root.join("unused.dump"); + let args = pull_args( + &["--dump", dump.to_str().unwrap()], + &root.join("project"), + ); + let stale = + stale_files(root, &BTreeMap::new(), &BTreeSet::new(), &args) + .unwrap(); + assert!(stale.is_empty(), "expected no stale files, got {stale:?}"); + } + + /// `--no-roles` also skips role extraction on a live connection, so + /// it must be treated the same as `--dump` + #[test] + fn no_roles_does_not_stale_roles_or_users() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + touch(root, "roles/app_group.yaml"); + touch(root, "users/app_login.yaml"); + let args = pull_args(&["--no-roles"], &root.join("project")); + let stale = + stale_files(root, &BTreeMap::new(), &BTreeSet::new(), &args) + .unwrap(); + assert!(stale.is_empty(), "expected no stale files, got {stale:?}"); + } + + /// Without `--dump`/`--no-roles`, a live pull does extract roles, so + /// an unmatched file under `roles/`/`users/` is still stale — the + /// protection above must not blanket-exempt those directories + #[test] + fn live_source_still_stales_roles_and_users() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + touch(root, "roles/app_group.yaml"); + touch(root, "users/app_login.yaml"); + let args = pull_args(&[], &root.join("project")); + let stale = + stale_files(root, &BTreeMap::new(), &BTreeSet::new(), &args) + .unwrap(); + assert_eq!( + stale, + vec![ + PathBuf::from("roles/app_group.yaml"), + PathBuf::from("users/app_login.yaml"), + ] + ); + } + + /// A schema excluded via `--exclude-schema` was deliberately left + /// out of the dump, so its files are not stale even though the + /// render did not produce them + #[test] + fn excluded_schema_is_not_staled() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + touch(root, "schemata/reporting.yaml"); + touch(root, "tables/reporting/big.yaml"); + let args = + pull_args(&["--exclude-schema", "report*"], &root.join("project")); + let stale = + stale_files(root, &BTreeMap::new(), &BTreeSet::new(), &args) + .unwrap(); + assert!(stale.is_empty(), "expected no stale files, got {stale:?}"); + } + + /// A table excluded via `--exclude-table` (schema-qualified or bare) + /// is not stale, but other tables in the same schema still are + #[test] + fn excluded_table_is_not_staled() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + touch(root, "tables/public/big.yaml"); + touch(root, "tables/public/small.yaml"); + let args = pull_args( + &["--exclude-table", "public.big"], + &root.join("project"), + ); + let stale = + stale_files(root, &BTreeMap::new(), &BTreeSet::new(), &args) + .unwrap(); + assert_eq!(stale, vec![PathBuf::from("tables/public/small.yaml")]); + } + + #[test] + fn glob_match_wildcards() { + assert!(glob_match("report*", "reporting")); + assert!(glob_match("*_vw", "orders_vw")); + assert!(glob_match("public.big", "public.big")); + assert!(!glob_match("public.big", "public.small")); + assert!(glob_match("pub??c.big", "public.big")); + } +} diff --git a/src/pull/writer.rs b/src/pull/writer.rs index 543b2d1..9e6094d 100644 --- a/src/pull/writer.rs +++ b/src/pull/writer.rs @@ -24,18 +24,15 @@ pub fn render( }; writer.write_project_file(assembly)?; for schema in &assembly.schemas { - writer.save( - Path::new("schemata").join(format!("{}.yaml", schema.name)), - schema, - )?; + writer.save(top_level("schemata", &schema.name)?, schema)?; } for domain in &assembly.domains { writer - .save(nested("domains", &domain.schema, &domain.name), domain)?; + .save(nested("domains", &domain.schema, &domain.name)?, domain)?; } for sequence in &assembly.sequences { writer.save( - nested("sequences", &sequence.schema, &sequence.name), + nested("sequences", &sequence.schema, &sequence.name)?, sequence, )?; } @@ -47,26 +44,23 @@ pub fn render( map.insert(String::from("dependencies"), dependencies); } writer.save_value( - nested("tables", &table.schema, &table.name), + nested("tables", &table.schema, &table.name)?, &value, )?; } for view in &assembly.views { - writer.save(nested("views", &view.schema, &view.name), view)?; + writer.save(nested("views", &view.schema, &view.name)?, view)?; } for view in &assembly.materialized_views { writer.save( - nested("materialized_views", &view.schema, &view.name), + nested("materialized_views", &view.schema, &view.name)?, view, )?; } writer.write_functions(assembly)?; writer.write_types(assembly)?; for server in &assembly.servers { - writer.save( - Path::new("servers").join(format!("{}.yaml", server.name)), - server, - )?; + writer.save(top_level("servers", &server.name)?, server)?; } writer.write_user_mappings(assembly)?; writer.write_roles(assembly)?; @@ -198,7 +192,7 @@ impl Writer { } used.insert((function.schema.clone(), filename.clone())); self.save( - nested("functions", &function.schema, &filename), + nested("functions", &function.schema, &filename)?, function, )?; } @@ -218,10 +212,7 @@ impl Writer { .filter(|t| t.schema == schema) .collect(); let container = json!({"schema": schema, "types": types}); - self.save_value( - Path::new("types").join(format!("{schema}.yaml")), - &container, - )?; + self.save_value(top_level("types", schema)?, &container)?; } Ok(()) } @@ -248,10 +239,7 @@ impl Writer { options: user_options(state), settings, }; - self.save( - Path::new("users").join(format!("{name}.yaml")), - &user, - )?; + self.save(top_level("users", name)?, &user)?; } else { let role = models::Role { name: name.clone(), @@ -264,10 +252,7 @@ impl Writer { .then(|| state.options.clone()), settings, }; - self.save( - Path::new("roles").join(format!("{name}.yaml")), - &role, - )?; + self.save(top_level("roles", name)?, &role)?; } } Ok(()) @@ -291,11 +276,7 @@ impl Writer { } } } - self.save( - Path::new("user_mappings") - .join(format!("{}.yaml", mapping.name)), - &mapping, - )?; + self.save(top_level("user_mappings", &mapping.name)?, &mapping)?; } Ok(()) } @@ -415,10 +396,39 @@ fn serialize(value: &T) -> Result { .map_err(|e| format!("failed to serialize: {e}")) } -fn nested(directory: &str, schema: &str, name: &str) -> PathBuf { - Path::new(directory) - .join(schema) - .join(format!("{name}.yaml")) +fn nested( + directory: &str, + schema: &str, + name: &str, +) -> Result { + Ok(Path::new(directory) + .join(safe_component(schema)?) + .join(format!("{}.yaml", safe_component(name)?))) +} + +/// A `directory/name.yaml` path with the name validated as a safe path +/// component +fn top_level(directory: &str, name: &str) -> Result { + Ok(Path::new(directory).join(format!("{}.yaml", safe_component(name)?))) +} + +/// Reject a database identifier that would escape or redirect the +/// destination when used as a filesystem path segment. PostgreSQL +/// quoted identifiers may contain `/`, `\`, or `.` sequences, and +/// `Path::join` treats an absolute or `..` component as a traversal. +fn safe_component(component: &str) -> Result<&str, String> { + if component.is_empty() + || component == "." + || component == ".." + || component.contains('/') + || component.contains('\\') + { + return Err(format!( + "refusing to write to a path derived from the unsafe database \ + identifier {component:?}" + )); + } + Ok(component) } pub(crate) fn read_ignore( @@ -494,3 +504,45 @@ fn walk_directories(root: &Path) -> Result, String> { } Ok(results) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn safe_component_accepts_normal_identifiers() { + assert_eq!(safe_component("public").unwrap(), "public"); + assert_eq!(safe_component("My Table").unwrap(), "My Table"); + assert_eq!(safe_component("weird.name").unwrap(), "weird.name"); + } + + #[test] + fn safe_component_rejects_traversal() { + for hostile in ["", ".", "..", "/etc", "../../tmp/x", "a/b", "a\\b"] { + assert!( + safe_component(hostile).is_err(), + "expected {hostile:?} to be rejected" + ); + } + } + + #[test] + fn nested_rejects_hostile_schema_or_name() { + assert!(nested("tables", "../../../tmp", "x").is_err()); + assert!(nested("tables", "public", "..").is_err()); + assert!(nested("tables", "/etc", "x").is_err()); + assert_eq!( + nested("tables", "public", "orders").unwrap(), + Path::new("tables").join("public").join("orders.yaml") + ); + } + + #[test] + fn top_level_rejects_hostile_name() { + assert!(top_level("roles", "../../etc/passwd").is_err()); + assert_eq!( + top_level("roles", "admin").unwrap(), + Path::new("roles").join("admin.yaml") + ); + } +} diff --git a/src/utils.rs b/src/utils.rs index 6834e3c..0a83241 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -2,13 +2,127 @@ use serde_json::Value; +/// PostgreSQL keywords that pg_dump's `fmtId` quotes even when they +/// otherwise look like a safe unquoted identifier. This covers the +/// `RESERVED_KEYWORD` and `TYPE_FUNC_NAME_KEYWORD` sets. The +/// `COL_NAME_KEYWORD` set (e.g. `int`, `timestamp`, `values`) is +/// deliberately excluded: those double as type names and +/// `quote_ident` is also applied to type names in some render paths +/// (e.g. CAST target types), where quoting them would diverge from +/// pg_dump. Must stay sorted for `binary_search`. +const RESERVED_KEYWORDS: &[&str] = &[ + "all", + "analyse", + "analyze", + "and", + "any", + "array", + "as", + "asc", + "asymmetric", + "authorization", + "binary", + "both", + "case", + "cast", + "check", + "collate", + "collation", + "column", + "concurrently", + "constraint", + "create", + "cross", + "current_catalog", + "current_date", + "current_role", + "current_schema", + "current_time", + "current_timestamp", + "current_user", + "default", + "deferrable", + "desc", + "distinct", + "do", + "else", + "end", + "except", + "false", + "fetch", + "for", + "foreign", + "freeze", + "from", + "full", + "grant", + "group", + "having", + "ilike", + "in", + "initially", + "inner", + "intersect", + "into", + "is", + "isnull", + "join", + "lateral", + "leading", + "left", + "like", + "limit", + "localtime", + "localtimestamp", + "natural", + "not", + "notnull", + "null", + "offset", + "on", + "only", + "or", + "order", + "outer", + "overlaps", + "placing", + "primary", + "references", + "returning", + "right", + "select", + "session_user", + "similar", + "some", + "symmetric", + "table", + "tablesample", + "then", + "to", + "trailing", + "true", + "union", + "unique", + "user", + "using", + "variadic", + "verbose", + "when", + "where", + "window", + "with", +]; + /// Quote a PostgreSQL identifier (object name, etc) pub fn quote_ident(value: &str) -> String { - if !value.is_empty() + let is_safe_shape = !value.is_empty() && value .chars() .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '_') - { + && !value.as_bytes()[0].is_ascii_digit(); + let is_reserved = + is_safe_shape && RESERVED_KEYWORDS.binary_search(&value).is_ok(); + if is_safe_shape && !is_reserved { value.to_string() } else { format!("\"{}\"", value.replace('"', "\"\"")) @@ -77,6 +191,21 @@ mod tests { assert_eq!(quote_ident("uuid-ossp"), "\"uuid-ossp\""); assert_eq!(quote_ident("==="), "\"===\""); assert_eq!(quote_ident("Has\"Quote"), "\"Has\"\"Quote\""); + assert_eq!(quote_ident("orders"), "orders"); + assert_eq!(quote_ident("my_table"), "my_table"); + assert_eq!(quote_ident("order"), "\"order\""); + assert_eq!(quote_ident("user"), "\"user\""); + assert_eq!(quote_ident("2fa"), "\"2fa\""); + // `all` is a RESERVED keyword pg_dump also quotes + assert_eq!(quote_ident("all"), "\"all\""); + } + + #[test] + fn reserved_keywords_stay_sorted() { + assert!( + RESERVED_KEYWORDS.windows(2).all(|w| w[0] < w[1]), + "RESERVED_KEYWORDS must be sorted and unique for binary_search" + ); } #[test] diff --git a/test-project/tables/test/users.yaml b/test-project/tables/test/users.yaml index 2325fe9..f50444c 100644 --- a/test-project/tables/test/users.yaml +++ b/test-project/tables/test/users.yaml @@ -14,9 +14,9 @@ columns: data_type: TIMESTAMP WITH TIME ZONE default: CURRENT_TIMESTAMP nullable: false - - name: last_modified_at, - comment: When the record was last modified, - data_type: TIMESTAMP WITH TIME ZONE, + - name: last_modified_at + comment: When the record was last modified + data_type: TIMESTAMP WITH TIME ZONE nullable: true - name: state comment: The current state of the user diff --git a/tests/build_parity.rs b/tests/build_parity.rs index 0fe969f..3e68d54 100644 --- a/tests/build_parity.rs +++ b/tests/build_parity.rs @@ -61,6 +61,11 @@ const DEVIATIONS: &[(&str, &str, &str)] = &[ "test", "Copied from the snowball template", ), + // Python rendered COMMENT ON CAST with an erroneous schema prefix + // (a cast has no namespace) and COMMENT ON OPERATOR without the + // argument signature required to disambiguate it + ("COMMENT", "test", "(int4 AS timestamp)"), + ("COMMENT", "test", "==="), ]; /// (desc, namespace, tag, required defn fragment) for the corrected @@ -145,6 +150,44 @@ const CORRECTED: &[(&str, &str, &str, &str)] = &[ "custom_snowball", "(INIT = dsnowball_init, LEXIZE = dsnowball_lexize)", ), + ( + "COMMENT", + "test", + "(int4 AS timestamp)", + "COMMENT ON CAST (int4 AS timestamp) IS", + ), + ( + "COMMENT", + "test", + "===", + "COMMENT ON OPERATOR test.=== (box, box) IS", + ), +]; + +/// (tag, comment) for COMMENT ON COLUMN entries: build now emits these +/// (data loss fix — Python's build.py never rendered column comments) +const NEW_COLUMN_COMMENTS: &[(&str, &str)] = &[ + ("addresses.created_at", "When the record was created"), + ("addresses.id", "The user ID"), + ( + "addresses.last_modified_at", + "When the record was last modified", + ), + ("addresses.user_id", "Foreign Key to test.users"), + ("empty_table.created_at", "When the record was created"), + ("empty_table.id", "The auto-incrementing row ID value"), + ( + "empty_table.last_modified_at", + "When the record was last modified", + ), + ("empty_table.value", "Some random value"), + ("users.created_at", "When the record was created"), + ("users.id", "The user ID"), + ( + "users.last_modified_at", + "When the record was last modified", + ), + ("users.state", "The current state of the user"), ]; /// Comment entries Python lost entirely to the text search name bug @@ -231,6 +274,9 @@ fn matches_python_build_output() { && !RECOVERED_COMMENTS .iter() .any(|(tag, _)| key == &entry_key("COMMENT", "test", tag)) + && !NEW_COLUMN_COMMENTS + .iter() + .any(|(tag, _)| key == &entry_key("COMMENT", "test", tag)) }) .cloned() .collect(); @@ -264,6 +310,17 @@ fn matches_python_build_output() { assert!(defn.contains(comment), "{key:?} missing comment text"); } + // column comments Python's build.py dropped entirely (data loss) + for (tag, comment) in NEW_COLUMN_COMMENTS { + let key = entry_key("COMMENT", "test", tag); + let defn = rust_tuples + .iter() + .find(|(k, ..)| k == &key) + .map(|(_, _, defn, ..)| defn.as_str()) + .unwrap_or_else(|| panic!("missing column comment {key:?}")); + assert!(defn.contains(comment), "{key:?} missing comment text"); + } + // the developers group's grant emits an ACL entry, which the // Python build never did let acl = rust_tuples @@ -277,9 +334,9 @@ fn matches_python_build_output() { public.empty_table TO developers;\n" ); - // 60 Python entries + 4 recovered text search comments + 1 ACL - // - 1 create: false role - assert_eq!(dump.entries().len(), 64); + // 60 Python entries + 4 recovered text search comments + 12 new + // column comments + 1 ACL - 1 create: false role + assert_eq!(dump.entries().len(), 76); } #[test]