Postgres: support WHERE index_predicate in the ON CONFLICT target

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;

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<Ident>),
+    /// 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 <https://www.postgresql.org/docs/current/sql-insert.html>
+    /// and <https://www.sqlite.org/lang_upsert.html>
+    Columns {
+        /// The columns of the conflict target.
+        columns: Vec<Ident>,
+        /// The `WHERE index_predicate` after the columns, if any.
+        predicate: Option<Expr>,
+    },
     /// 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}"),
         }
     }

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#"

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
                         };

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,
                 })),
             ..

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 <https://www.sqlite.org/lang_upsert.html>
+    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 {})])
 }