diff --git a/fuzz/fuzz_seeds/postgres_on_conflict_target_predicate.sql b/fuzz/fuzz_seeds/postgres_on_conflict_target_predicate.sql new file mode 100644 index 0000000000..a63d004f48 --- /dev/null +++ b/fuzz/fuzz_seeds/postgres_on_conflict_target_predicate.sql @@ -0,0 +1,2 @@ +INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT (a) WHERE b > 0 DO NOTHING; +INSERT INTO t (a, b, c) VALUES (2, 2, 3) ON CONFLICT (a, c) WHERE b > 0 AND c IS NOT NULL DO UPDATE SET b = EXCLUDED.b WHERE t.b < 10; diff --git a/src/ast/mod.rs b/src/ast/mod.rs index 1f2afd7fec..1c93c35d4a 100644 --- a/src/ast/mod.rs +++ b/src/ast/mod.rs @@ -6917,8 +6917,18 @@ pub struct OnConflict { #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))] /// Target specification for an `ON CONFLICT` clause. pub enum ConflictTarget { - /// Target specified as a list of columns. - Columns(Vec), + /// Target specified as a list of columns, optionally followed by + /// `WHERE index_predicate`, which lets a partial unique index be inferred as the arbiter: + /// `ON CONFLICT (a) WHERE b > 0 DO ...`. PostgreSQL and SQLite both accept the predicate. + /// + /// See + /// and + Columns { + /// The columns of the conflict target. + columns: Vec, + /// The `WHERE index_predicate` after the columns, if any. + predicate: Option, + }, /// Target specified as a named constraint. OnConstraint(ObjectName), } @@ -6968,7 +6978,13 @@ impl fmt::Display for OnConflict { impl fmt::Display for ConflictTarget { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - ConflictTarget::Columns(cols) => write!(f, "({})", display_comma_separated(cols)), + ConflictTarget::Columns { columns, predicate } => { + write!(f, "({})", display_comma_separated(columns))?; + if let Some(predicate) = predicate { + write!(f, " WHERE {predicate}")?; + } + Ok(()) + } ConflictTarget::OnConstraint(name) => write!(f, " ON CONSTRAINT {name}"), } } diff --git a/src/ast/spans.rs b/src/ast/spans.rs index 34294e4674..06eed733cf 100644 --- a/src/ast/spans.rs +++ b/src/ast/spans.rs @@ -1428,7 +1428,12 @@ impl Spanned for OnConflict { impl Spanned for ConflictTarget { fn span(&self) -> Span { match self { - ConflictTarget::Columns(vec) => union_spans(vec.iter().map(|i| i.span)), + ConflictTarget::Columns { columns, predicate } => union_spans( + columns + .iter() + .map(|i| i.span) + .chain(predicate.iter().map(|p| p.span())), + ), ConflictTarget::OnConstraint(object_name) => object_name.span(), } } @@ -2876,6 +2881,31 @@ ALTER TABLE users assert_eq!(stmt_span.end, (4, 12).into()); } + #[test] + fn test_on_conflict_target_predicate_span() { + let dialect = crate::dialect::PostgreSqlDialect {}; + let conflict_span = |sql: &str| -> Span { + let stmts = Parser::parse_sql(&dialect, sql).unwrap(); + let Statement::Insert(Insert { + on: Some(OnInsert::OnConflict(on_conflict)), + .. + }) = &stmts[0] + else { + panic!("not an INSERT ... ON CONFLICT statement"); + }; + on_conflict.span() + }; + + // The target's predicate extends the clause's span past the column list. + let sql = "INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT (a) WHERE b > 0 DO NOTHING"; + let test = SpanTest::new(&dialect, sql); + assert_eq!(test.get_source(conflict_span(sql)), "a) WHERE b > 0"); + + let sql = "INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT (a) DO NOTHING"; + let test = SpanTest::new(&dialect, sql); + assert_eq!(test.get_source(conflict_span(sql)), "a"); + } + #[test] fn test_replace_statement_span() { let sql = r#" diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 188d0a26ea..c7f565586e 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -18958,9 +18958,15 @@ impl<'a> Parser<'a> { if self.parse_keywords(&[Keyword::ON, Keyword::CONSTRAINT]) { Some(ConflictTarget::OnConstraint(self.parse_object_name(false)?)) } else if self.peek_token_ref().token == Token::LParen { - Some(ConflictTarget::Columns( - self.parse_parenthesized_column_list(IsOptional::Mandatory, false)?, - )) + let columns = + self.parse_parenthesized_column_list(IsOptional::Mandatory, false)?; + // PostgreSQL: `( column [, ...] ) [ WHERE index_predicate ]`. + let predicate = if self.parse_keyword(Keyword::WHERE) { + Some(self.parse_expr()?) + } else { + None + }; + Some(ConflictTarget::Columns { columns, predicate }) } else { None }; diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs index ac6d69f542..761d44d060 100644 --- a/tests/sqlparser_postgres.rs +++ b/tests/sqlparser_postgres.rs @@ -2469,6 +2469,76 @@ fn parse_prepare() { ); } +#[test] +fn parse_pg_on_conflict_target_predicate() { + for (sql, action_selection) in [ + ( + "INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT(a) WHERE b > 0 DO NOTHING", + None, + ), + ( + "INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT(a) WHERE b > 0 DO UPDATE SET b = EXCLUDED.b", + None, + ), + // The target's predicate and the action's condition are separate clauses. + ( + "INSERT INTO t (a, b) VALUES (1, 2) ON CONFLICT(a, c) WHERE b > 0 DO UPDATE SET b = EXCLUDED.b WHERE t.b < 10", + Some("t.b < 10"), + ), + ] { + match pg_and_generic().verified_stmt(sql) { + Statement::Insert(Insert { + on: + Some(OnInsert::OnConflict(OnConflict { + conflict_target: + Some(ConflictTarget::Columns { + predicate: Some(predicate), + .. + }), + action, + })), + .. + }) => { + assert_eq!(predicate.to_string(), "b > 0"); + let selection = match action { + OnConflictAction::DoUpdate(DoUpdate { selection, .. }) => selection, + OnConflictAction::DoNothing => None, + }; + assert_eq!(selection.map(|e| e.to_string()).as_deref(), action_selection); + } + statement => panic!("Expected INSERT statement, got: {statement:?}"), + } + } + + // A predicate only follows a column list: not `ON CONSTRAINT`, not an omitted target. + for sql in [ + "INSERT INTO t (a) VALUES (1) ON CONFLICT ON CONSTRAINT t_a WHERE a > 0 DO NOTHING", + "INSERT INTO t (a) VALUES (1) ON CONFLICT WHERE a > 0 DO NOTHING", + ] { + assert!( + pg_and_generic().parse_sql_statements(sql).is_err(), + "{sql} should not parse" + ); + } + + // The predicate is an expression, and the action must still follow it. + for (sql, err) in [ + ( + "INSERT INTO t (a) VALUES (1) ON CONFLICT (a) WHERE", + "Expected: an expression, found: EOF", + ), + ( + "INSERT INTO t (a) VALUES (1) ON CONFLICT (a) WHERE a > 0", + "Expected: DO, found: EOF", + ), + ] { + assert_eq!( + pg_and_generic().parse_sql_statements(sql).unwrap_err(), + ParserError::ParserError(err.to_string()), + ); + } +} + #[test] fn parse_pg_on_conflict() { let stmt = pg_and_generic().verified_stmt( @@ -2481,7 +2551,11 @@ fn parse_pg_on_conflict() { Statement::Insert(Insert { on: Some(OnInsert::OnConflict(OnConflict { - conflict_target: Some(ConflictTarget::Columns(cols)), + conflict_target: + Some(ConflictTarget::Columns { + columns: cols, + predicate: None, + }), action, })), .. @@ -2513,7 +2587,11 @@ fn parse_pg_on_conflict() { Statement::Insert(Insert { on: Some(OnInsert::OnConflict(OnConflict { - conflict_target: Some(ConflictTarget::Columns(cols)), + conflict_target: + Some(ConflictTarget::Columns { + columns: cols, + predicate: None, + }), action, })), .. @@ -2575,7 +2653,11 @@ fn parse_pg_on_conflict() { Statement::Insert(Insert { on: Some(OnInsert::OnConflict(OnConflict { - conflict_target: Some(ConflictTarget::Columns(cols)), + conflict_target: + Some(ConflictTarget::Columns { + columns: cols, + predicate: None, + }), action, })), .. diff --git a/tests/sqlparser_sqlite.rs b/tests/sqlparser_sqlite.rs index 07a279ae5d..fc2a73bcc8 100644 --- a/tests/sqlparser_sqlite.rs +++ b/tests/sqlparser_sqlite.rs @@ -1362,6 +1362,31 @@ fn parse_over_as_column_alias_after_function() { ); } +#[test] +fn parse_on_conflict_target_predicate() { + // SQLite's upsert target takes the same `( indexed-column, ... ) WHERE expr` as PostgreSQL's. + // See + let sql = "INSERT INTO t VALUES (1, 2, 3) ON CONFLICT(a) WHERE b > 0 DO NOTHING"; + match sqlite().verified_stmt(sql) { + Statement::Insert(Insert { + on: + Some(OnInsert::OnConflict(OnConflict { + conflict_target: + Some(ConflictTarget::Columns { + columns, + predicate: Some(predicate), + }), + action: OnConflictAction::DoNothing, + })), + .. + }) => { + assert_eq!(columns, vec![Ident::new("a")]); + assert_eq!(predicate.to_string(), "b > 0"); + } + other => panic!("unexpected statement: {other:?}"), + } +} + fn sqlite() -> TestedDialects { TestedDialects::new(vec![Box::new(SQLiteDialect {})]) }