Skip to content

Commit 4b24f9e

Browse files
fmguerreirofmguerreiro
authored
PostgreSQL: extend COMMENT ON for trigger/aggregate/policy and dollar-quoted bodies (#25)
* PostgreSQL: extend COMMENT ON for trigger/aggregate/policy and dollar-quoted bodies Adds Trigger, Aggregate, and Policy variants to CommentObject and extends Statement::Comment with optional `arguments: Option<Vec<DataType>>` (for FUNCTION/PROCEDURE/AGGREGATE overload signatures) and `relation: Option<ObjectName>` (for the `ON <table>` tail of TRIGGER and POLICY). Parser now accepts: - COMMENT ON AGGREGATE name(arg, ...) IS '...' - COMMENT ON TRIGGER name ON table IS '...' - COMMENT ON POLICY name ON table IS '...' - COMMENT ON FUNCTION/PROCEDURE name(arg, ...) IS '...' parse_literal_string now accepts Token::DollarQuotedString in the PostgreSQL/Generic dialects, enabling COMMENT ON ... IS $$body$$ and $tag$body$tag$. NOTE: CommentObject and Statement::Comment are not #[non_exhaustive], so new variants and new struct fields are a breaking change for downstream consumers that pattern-match exhaustively. Recommend bumping to 0.61.0 for the next release rather than a 0.60.x patch. * PostgreSQL: require argument list for COMMENT ON AGGREGATE, escape quotes in Display - Reject bare `COMMENT ON AGGREGATE foo IS '...'`; PostgreSQL requires an argument list as part of the aggregate's identity (distinguishes ordinary vs ordered-set forms). FUNCTION and PROCEDURE remain tolerant of the bare form for backward compatibility. - Escape embedded single quotes in the Statement::Comment Display impl via value::escape_single_quote_string. This was latent but became reachable with dollar-quoted parsing (e.g. $$it's$$ previously round-tripped to invalid SQL `'it's'`). - Add negative parser test for the AGGREGATE guard and a round-trip test for a dollar-quoted body containing a single quote.
1 parent 759ab21 commit 4b24f9e

4 files changed

Lines changed: 232 additions & 2 deletions

File tree

‎src/ast/mod.rs‎

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2460,6 +2460,8 @@ impl fmt::Display for ShowCreateObject {
24602460
#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
24612461
/// Objects that can be targeted by a `COMMENT` statement.
24622462
pub enum CommentObject {
2463+
/// An aggregate function.
2464+
Aggregate,
24632465
/// A collation.
24642466
Collation,
24652467
/// A table column.
@@ -2476,6 +2478,8 @@ pub enum CommentObject {
24762478
Index,
24772479
/// A materialized view.
24782480
MaterializedView,
2481+
/// A row-level security policy.
2482+
Policy,
24792483
/// A procedure.
24802484
Procedure,
24812485
/// A role.
@@ -2486,6 +2490,8 @@ pub enum CommentObject {
24862490
Sequence,
24872491
/// A table.
24882492
Table,
2493+
/// A trigger.
2494+
Trigger,
24892495
/// A type.
24902496
Type,
24912497
/// A user.
@@ -2497,6 +2503,7 @@ pub enum CommentObject {
24972503
impl fmt::Display for CommentObject {
24982504
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
24992505
match self {
2506+
CommentObject::Aggregate => f.write_str("AGGREGATE"),
25002507
CommentObject::Collation => f.write_str("COLLATION"),
25012508
CommentObject::Column => f.write_str("COLUMN"),
25022509
CommentObject::Database => f.write_str("DATABASE"),
@@ -2505,11 +2512,13 @@ impl fmt::Display for CommentObject {
25052512
CommentObject::Function => f.write_str("FUNCTION"),
25062513
CommentObject::Index => f.write_str("INDEX"),
25072514
CommentObject::MaterializedView => f.write_str("MATERIALIZED VIEW"),
2515+
CommentObject::Policy => f.write_str("POLICY"),
25082516
CommentObject::Procedure => f.write_str("PROCEDURE"),
25092517
CommentObject::Role => f.write_str("ROLE"),
25102518
CommentObject::Schema => f.write_str("SCHEMA"),
25112519
CommentObject::Sequence => f.write_str("SEQUENCE"),
25122520
CommentObject::Table => f.write_str("TABLE"),
2521+
CommentObject::Trigger => f.write_str("TRIGGER"),
25132522
CommentObject::Type => f.write_str("TYPE"),
25142523
CommentObject::User => f.write_str("USER"),
25152524
CommentObject::View => f.write_str("VIEW"),
@@ -4404,6 +4413,15 @@ pub enum Statement {
44044413
object_type: CommentObject,
44054414
/// Name of the object the comment applies to.
44064415
object_name: ObjectName,
4416+
/// Argument types for overloaded objects. `Some(vec![])` represents an
4417+
/// empty `()` parameter list (e.g. `COMMENT ON FUNCTION f() IS '…'`),
4418+
/// while `None` means no parameter list was provided. Used for
4419+
/// `FUNCTION`, `PROCEDURE`, and `AGGREGATE` targets.
4420+
arguments: Option<Vec<DataType>>,
4421+
/// Partner relation for objects scoped to a table, i.e. the
4422+
/// `ON <relation>` tail in `COMMENT ON TRIGGER t ON tbl IS '…'` or
4423+
/// `COMMENT ON POLICY p ON tbl IS '…'`.
4424+
relation: Option<ObjectName>,
44074425
/// Optional comment text (None to remove comment).
44084426
comment: Option<String>,
44094427
/// An optional `IF EXISTS` clause. (Non-standard.)
@@ -6187,16 +6205,25 @@ impl fmt::Display for Statement {
61876205
Statement::Comment {
61886206
object_type,
61896207
object_name,
6208+
arguments,
6209+
relation,
61906210
comment,
61916211
if_exists,
61926212
} => {
61936213
write!(f, "COMMENT ")?;
61946214
if *if_exists {
61956215
write!(f, "IF EXISTS ")?
61966216
};
6197-
write!(f, "ON {object_type} {object_name} IS ")?;
6217+
write!(f, "ON {object_type} {object_name}")?;
6218+
if let Some(args) = arguments {
6219+
write!(f, "({})", display_comma_separated(args))?;
6220+
}
6221+
if let Some(relation) = relation {
6222+
write!(f, " ON {relation}")?;
6223+
}
6224+
write!(f, " IS ")?;
61986225
if let Some(c) = comment {
6199-
write!(f, "'{c}'")
6226+
write!(f, "'{}'", value::escape_single_quote_string(c))
62006227
} else {
62016228
write!(f, "NULL")
62026229
}

‎src/parser/mod.rs‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -903,6 +903,9 @@ impl<'a> Parser<'a> {
903903
let token = self.next_token();
904904

905905
let (object_type, object_name) = match token.token {
906+
Token::Word(w) if w.keyword == Keyword::AGGREGATE => {
907+
(CommentObject::Aggregate, self.parse_object_name(false)?)
908+
}
906909
Token::Word(w) if w.keyword == Keyword::COLLATION => {
907910
(CommentObject::Collation, self.parse_object_name(false)?)
908911
}
@@ -931,6 +934,9 @@ impl<'a> Parser<'a> {
931934
self.parse_object_name(false)?,
932935
)
933936
}
937+
Token::Word(w) if w.keyword == Keyword::POLICY => {
938+
(CommentObject::Policy, self.parse_object_name(false)?)
939+
}
934940
Token::Word(w) if w.keyword == Keyword::PROCEDURE => {
935941
(CommentObject::Procedure, self.parse_object_name(false)?)
936942
}
@@ -946,6 +952,9 @@ impl<'a> Parser<'a> {
946952
Token::Word(w) if w.keyword == Keyword::TABLE => {
947953
(CommentObject::Table, self.parse_object_name(false)?)
948954
}
955+
Token::Word(w) if w.keyword == Keyword::TRIGGER => {
956+
(CommentObject::Trigger, self.parse_object_name(false)?)
957+
}
949958
Token::Word(w) if w.keyword == Keyword::TYPE => {
950959
(CommentObject::Type, self.parse_object_name(false)?)
951960
}
@@ -958,6 +967,33 @@ impl<'a> Parser<'a> {
958967
_ => self.expected("comment object_type", token)?,
959968
};
960969

970+
let arguments = match object_type {
971+
CommentObject::Function | CommentObject::Procedure | CommentObject::Aggregate => {
972+
if self.consume_token(&Token::LParen) {
973+
let list = self.parse_comma_separated0(Self::parse_data_type, Token::RParen)?;
974+
self.expect_token(&Token::RParen)?;
975+
Some(list)
976+
} else {
977+
None
978+
}
979+
}
980+
_ => None,
981+
};
982+
983+
if object_type == CommentObject::Aggregate && arguments.is_none() {
984+
return Err(ParserError::ParserError(
985+
"COMMENT ON AGGREGATE requires an argument list, e.g. AGGREGATE foo(int)".into(),
986+
));
987+
}
988+
989+
let relation = match object_type {
990+
CommentObject::Trigger | CommentObject::Policy => {
991+
self.expect_keyword_is(Keyword::ON)?;
992+
Some(self.parse_object_name(false)?)
993+
}
994+
_ => None,
995+
};
996+
961997
self.expect_keyword_is(Keyword::IS)?;
962998
let comment = if self.parse_keyword(Keyword::NULL) {
963999
None
@@ -967,6 +1003,8 @@ impl<'a> Parser<'a> {
9671003
Ok(Statement::Comment {
9681004
object_type,
9691005
object_name,
1006+
arguments,
1007+
relation,
9701008
comment,
9711009
if_exists,
9721010
})
@@ -12751,6 +12789,9 @@ impl<'a> Parser<'a> {
1275112789
Ok(s)
1275212790
}
1275312791
Token::UnicodeStringLiteral(s) => Ok(s),
12792+
Token::DollarQuotedString(s) if dialect_of!(self is PostgreSqlDialect | GenericDialect) => {
12793+
Ok(s.value)
12794+
}
1275412795
_ => self.expected("literal string", next_token),
1275512796
}
1275612797
}

‎tests/sqlparser_common.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15348,6 +15348,7 @@ fn parse_comments() {
1534815348
object_name,
1534915349
comment: Some(comment),
1535015350
if_exists,
15351+
..
1535115352
} => {
1535215353
assert_eq!("comment", comment);
1535315354
assert_eq!("tab.name", object_name.to_string());
@@ -15385,6 +15386,7 @@ fn parse_comments() {
1538515386
object_name,
1538615387
comment: Some(comment),
1538715388
if_exists,
15389+
..
1538815390
} => {
1538915391
assert_eq!("comment", comment);
1539015392
assert_eq!("db.t0", object_name.to_string());
@@ -15403,6 +15405,7 @@ fn parse_comments() {
1540315405
object_name,
1540415406
comment: None,
1540515407
if_exists,
15408+
..
1540615409
} => {
1540715410
assert_eq!("public.tab", object_name.to_string());
1540815411
assert_eq!(CommentObject::Table, object_type);

‎tests/sqlparser_postgres.rs‎

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1114,6 +1114,8 @@ fn parse_drop_and_comment_collation_ast() {
11141114
Statement::Comment {
11151115
object_type: CommentObject::Collation,
11161116
object_name: ObjectName::from(vec![Ident::new("test0")]),
1117+
arguments: None,
1118+
relation: None,
11171119
comment: Some("US English".to_string()),
11181120
if_exists: false,
11191121
}
@@ -10509,3 +10511,160 @@ fn alter_table_detach_partition_finalize() {
1050910511
_ => panic!("Expected AlterTable"),
1051010512
}
1051110513
}
10514+
10515+
#[test]
10516+
fn parse_comment_on_trigger() {
10517+
match pg_and_generic()
10518+
.verified_stmt("COMMENT ON TRIGGER my_trigger ON public.my_table IS 'trigger note'")
10519+
{
10520+
Statement::Comment {
10521+
object_type,
10522+
object_name,
10523+
relation,
10524+
arguments,
10525+
comment,
10526+
if_exists,
10527+
} => {
10528+
assert_eq!(CommentObject::Trigger, object_type);
10529+
assert_eq!("my_trigger", object_name.to_string());
10530+
assert_eq!("public.my_table", relation.unwrap().to_string());
10531+
assert!(arguments.is_none());
10532+
assert_eq!(Some("trigger note".to_string()), comment);
10533+
assert!(!if_exists);
10534+
}
10535+
_ => panic!("Expected COMMENT ON TRIGGER"),
10536+
}
10537+
10538+
pg_and_generic().verified_stmt("COMMENT ON TRIGGER t ON s.foo IS NULL");
10539+
pg_and_generic().verified_stmt("COMMENT IF EXISTS ON TRIGGER t ON foo IS 'note'");
10540+
}
10541+
10542+
#[test]
10543+
fn parse_comment_on_policy() {
10544+
match pg_and_generic()
10545+
.verified_stmt("COMMENT ON POLICY tenant_isolation ON public.docs IS 'rls'")
10546+
{
10547+
Statement::Comment {
10548+
object_type,
10549+
object_name,
10550+
relation,
10551+
arguments,
10552+
comment,
10553+
if_exists,
10554+
} => {
10555+
assert_eq!(CommentObject::Policy, object_type);
10556+
assert_eq!("tenant_isolation", object_name.to_string());
10557+
assert_eq!("public.docs", relation.unwrap().to_string());
10558+
assert!(arguments.is_none());
10559+
assert_eq!(Some("rls".to_string()), comment);
10560+
assert!(!if_exists);
10561+
}
10562+
_ => panic!("Expected COMMENT ON POLICY"),
10563+
}
10564+
10565+
pg_and_generic().verified_stmt("COMMENT ON POLICY p ON t IS NULL");
10566+
}
10567+
10568+
#[test]
10569+
fn parse_comment_on_aggregate() {
10570+
match pg_and_generic().verified_stmt("COMMENT ON AGGREGATE my_sum(INTEGER) IS 'sums things'") {
10571+
Statement::Comment {
10572+
object_type,
10573+
object_name,
10574+
relation,
10575+
arguments,
10576+
comment,
10577+
if_exists,
10578+
} => {
10579+
assert_eq!(CommentObject::Aggregate, object_type);
10580+
assert_eq!("my_sum", object_name.to_string());
10581+
assert!(relation.is_none());
10582+
let args = arguments.expect("aggregate should carry argument types");
10583+
assert_eq!(1, args.len());
10584+
assert!(matches!(args[0], DataType::Integer(_)));
10585+
assert_eq!(Some("sums things".to_string()), comment);
10586+
assert!(!if_exists);
10587+
}
10588+
_ => panic!("Expected COMMENT ON AGGREGATE"),
10589+
}
10590+
10591+
pg_and_generic().verified_stmt("COMMENT ON AGGREGATE public.avg(DOUBLE PRECISION) IS 'avg'");
10592+
pg_and_generic().verified_stmt("COMMENT ON AGGREGATE s.f(INTEGER, TEXT) IS NULL");
10593+
pg_and_generic().verified_stmt("COMMENT ON AGGREGATE no_args() IS 'none'");
10594+
}
10595+
10596+
#[test]
10597+
fn parse_comment_on_aggregate_requires_argument_list() {
10598+
let result = pg().parse_sql_statements("COMMENT ON AGGREGATE foo IS 'bad'");
10599+
assert!(
10600+
result.is_err(),
10601+
"COMMENT ON AGGREGATE without argument list must error, got: {result:?}"
10602+
);
10603+
}
10604+
10605+
#[test]
10606+
fn parse_comment_on_function_with_arg_types() {
10607+
match pg_and_generic().verified_stmt("COMMENT ON FUNCTION add(INTEGER, INTEGER) IS 'adds'") {
10608+
Statement::Comment {
10609+
object_type,
10610+
object_name,
10611+
relation,
10612+
arguments,
10613+
comment,
10614+
if_exists,
10615+
} => {
10616+
assert_eq!(CommentObject::Function, object_type);
10617+
assert_eq!("add", object_name.to_string());
10618+
assert!(relation.is_none());
10619+
let args = arguments.expect("function should carry argument types");
10620+
assert_eq!(2, args.len());
10621+
assert_eq!(Some("adds".to_string()), comment);
10622+
assert!(!if_exists);
10623+
}
10624+
_ => panic!("Expected COMMENT ON FUNCTION"),
10625+
}
10626+
10627+
pg_and_generic().verified_stmt("COMMENT ON FUNCTION public.noop() IS 'no args'");
10628+
pg_and_generic().verified_stmt("COMMENT ON FUNCTION f(TEXT) IS NULL");
10629+
match pg_and_generic().verified_stmt("COMMENT ON FUNCTION f IS 'bare'") {
10630+
Statement::Comment {
10631+
object_type,
10632+
arguments,
10633+
..
10634+
} => {
10635+
assert_eq!(CommentObject::Function, object_type);
10636+
assert!(arguments.is_none());
10637+
}
10638+
_ => panic!("Expected COMMENT ON FUNCTION"),
10639+
}
10640+
}
10641+
10642+
#[test]
10643+
fn parse_comment_dollar_quoted_body() {
10644+
pg_and_generic().one_statement_parses_to(
10645+
"COMMENT ON TABLE foo IS $$hello world$$",
10646+
"COMMENT ON TABLE foo IS 'hello world'",
10647+
);
10648+
10649+
pg_and_generic().one_statement_parses_to(
10650+
"COMMENT ON TABLE foo IS $tag$multi\nline$tag$",
10651+
"COMMENT ON TABLE foo IS 'multi\nline'",
10652+
);
10653+
10654+
pg_and_generic().one_statement_parses_to(
10655+
"COMMENT ON TABLE foo IS $$it's escaped$$",
10656+
"COMMENT ON TABLE foo IS 'it''s escaped'",
10657+
);
10658+
10659+
match pg_and_generic()
10660+
.parse_sql_statements("COMMENT ON TABLE foo IS $$hello$$")
10661+
.unwrap()
10662+
.pop()
10663+
.unwrap()
10664+
{
10665+
Statement::Comment { comment, .. } => {
10666+
assert_eq!(Some("hello".to_string()), comment);
10667+
}
10668+
_ => panic!("Expected COMMENT"),
10669+
}
10670+
}

0 commit comments

Comments
 (0)