Oracle: Support for MERGE predicates

Author: xitep

src/ast/mod.rs

@@ -8626,6 +8626,7 @@ impl Display for MergeInsertKind {
 ///
 /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/merge)
 /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement)
+/// [Oracle](https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/MERGE.html)
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -8644,14 +8645,22 @@ pub struct MergeInsertExpr {
     pub kind_token: AttachedToken,
     /// The insert type used by the statement.
     pub kind: MergeInsertKind,
+    /// An optional condition to restrict the insertion (Oracle specific)
+    ///
+    /// Enabled via [`Dialect::supports_merge_insert_predicate`](crate::dialect::Dialect::supports_merge_insert_predicate).
+    pub insert_predicate: Option<Expr>,
 }
 
 impl Display for MergeInsertExpr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if !self.columns.is_empty() {
             write!(f, "({}) ", display_comma_separated(self.columns.as_slice()))?;
         }
-        write!(f, "{}", self.kind)
+        write!(f, "{}", self.kind)?;
+        if let Some(predicate) = self.insert_predicate.as_ref() {
+            write!(f, " WHERE {}", predicate)?;
+        }
+        Ok(())
     }
 }
 
@@ -8664,6 +8673,7 @@ impl Display for MergeInsertExpr {
 ///
 /// [Snowflake](https://docs.snowflake.com/en/sql-reference/sql/merge)
 /// [BigQuery](https://cloud.google.com/bigquery/docs/reference/standard-sql/dml-syntax#merge_statement)
+/// [Oracle](https://docs.oracle.com/en/database/oracle/oracle-database/21/sqlrf/MERGE.html)
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
@@ -8684,7 +8694,16 @@ pub enum MergeAction {
     Update {
         /// The `UPDATE` token that starts the sub-expression.
         update_token: AttachedToken,
+        /// The update assiment expressions
         assignments: Vec<Assignment>,
+        /// `where_clause` for the update (Oralce specific)
+        ///
+        /// Enabled via [`Dialect::supports_merge_update_predicate`](crate::dialect::Dialect::supports_merge_update_predicate).
+        update_predicate: Option<Expr>,
+        /// `delete_clause` for the update "delete where" (Oracle specific)
+        ///
+        /// Enabled via [`Dialect::supports_merge_update_delete_predicate`](crate::dialect::Dialect::supports_merge_update_delete_predicate).
+        delete_predicate: Option<Expr>,
     },
     /// A plain `DELETE` clause
     Delete {
@@ -8699,8 +8718,20 @@ impl Display for MergeAction {
             MergeAction::Insert(insert) => {
                 write!(f, "INSERT {insert}")
             }
-            MergeAction::Update { assignments, .. } => {
-                write!(f, "UPDATE SET {}", display_comma_separated(assignments))
+            MergeAction::Update {
+                update_token: _,
+                assignments,
+                update_predicate,
+                delete_predicate,
+            } => {
+                write!(f, "UPDATE SET {}", display_comma_separated(assignments))?;
+                if let Some(predicate) = update_predicate.as_ref() {
+                    write!(f, " WHERE {predicate}")?;
+                }
+                if let Some(predicate) = delete_predicate.as_ref() {
+                    write!(f, " DELETE WHERE {predicate}")?;
+                }
+                Ok(())
             }
             MergeAction::Delete { .. } => {
                 write!(f, "DELETE")

src/ast/spans.rs

@@ -2420,8 +2420,13 @@ impl Spanned for MergeAction {
             MergeAction::Update {
                 update_token,
                 assignments,
+                update_predicate,
+                delete_predicate,
             } => union_spans(
-                core::iter::once(update_token.0.span).chain(assignments.iter().map(Spanned::span)),
+                core::iter::once(update_token.0.span)
+                    .chain(assignments.iter().map(Spanned::span))
+                    .chain(update_predicate.iter().map(Spanned::span))
+                    .chain(delete_predicate.iter().map(Spanned::span)),
             ),
             MergeAction::Delete { delete_token } => delete_token.0.span,
         }
@@ -2440,6 +2445,7 @@ impl Spanned for MergeInsertExpr {
                 },
             ]
             .into_iter()
+            .chain(self.insert_predicate.iter().map(Spanned::span))
             .chain(self.columns.iter().map(|i| i.span)),
         )
     }
@@ -2813,6 +2819,8 @@ WHERE id = 1
         if let MergeAction::Update {
             update_token,
             assignments: _,
+            update_predicate: _,
+            delete_predicate: _,
         } = &clauses[1].action
         {
             assert_eq!(
@@ -2933,4 +2941,44 @@ WHERE id = 1
             panic!("not a MERGE statement");
         };
     }
+
+    #[test]
+    fn test_merge_statement_spans_with_update_predicates() {
+        let sql = r#"
+       MERGE INTO a USING b ON a.id = b.id
+        WHEN MATCHED THEN
+              UPDATE set a.x = a.x + b.x
+               WHERE b.x != 2
+              DELETE WHERE a.x <> 3"#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        // ~ assert the span of the whole statement
+        let stmt_span = r[0].span();
+        assert_eq!(
+            stmt_span,
+            Span::new(Location::new(2, 8), Location::new(6, 36))
+        );
+    }
+
+    #[test]
+    fn test_merge_statement_spans_with_insert_predicate() {
+        let sql = r#"
+       MERGE INTO a USING b ON a.id = b.id
+        WHEN NOT MATCHED THEN
+            INSERT VALUES (b.x, b.y) WHERE b.x != 2
+-- qed
+"#;
+
+        let r = Parser::parse_sql(&crate::dialect::GenericDialect, sql).unwrap();
+        assert_eq!(1, r.len());
+
+        // ~ assert the span of the whole statement
+        let stmt_span = r[0].span();
+        assert_eq!(
+            stmt_span,
+            Span::new(Location::new(2, 8), Location::new(4, 52))
+        );
+    }
 }

src/dialect/generic.rs

@@ -195,4 +195,20 @@ impl Dialect for GenericDialect {
     fn supports_interval_options(&self) -> bool {
         true
     }
+
+    fn supports_merge_insert_qualified_columns(&self) -> bool {
+        true
+    }
+
+    fn supports_merge_insert_predicate(&self) -> bool {
+        true
+    }
+
+    fn supports_merge_update_predicate(&self) -> bool {
+        true
+    }
+
+    fn supports_merge_update_delete_predicate(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -601,13 +601,122 @@ pub trait Dialect: Debug + Any {
         false
     }
 
-    /// Return true if the dialect supports specifying multiple options
+    /// Returns true if the dialect supports specifying multiple options
     /// in a `CREATE TABLE` statement for the structure of the new table. For example:
     /// `CREATE TABLE t (a INT, b INT) AS SELECT 1 AS b, 2 AS a`
     fn supports_create_table_multi_schema_info_sources(&self) -> bool {
         false
     }
 
+    /// Returns `true` if the dialect supports qualified column names
+    /// as part of a MERGE's INSERT's column list. Example:
+    ///
+    /// ```sql
+    /// MERGE INTO FOO
+    /// USING FOO_IMP
+    ///    ON (FOO.ID = FOO_IMP.ID)
+    ///  WHEN NOT MATCHED THEN
+    ///      -- no qualifier
+    ///      INSERT (ID, NAME)
+    ///      VALUES (FOO_IMP.ID, UPPER(FOO_IMP.NAME))
+    /// ```
+    /// vs.
+    /// ```sql
+    /// MERGE INTO FOO
+    /// USING FOO_IMP
+    ///    ON (FOO.ID = FOO_IMP.ID)
+    ///  WHEN NOT MATCHED THEN
+    ///      -- here: qualified
+    ///      INSERT (FOO.ID, FOO.NAME)
+    ///      VALUES (FOO_IMP.ID, UPPER(FOO_IMP.NAME))
+    /// ```
+    /// or
+    /// ```sql
+    /// MERGE INTO FOO X
+    /// USING FOO_IMP
+    ///    ON (X.ID = FOO_IMP.ID)
+    ///  WHEN NOT MATCHED THEN
+    ///      -- here: qualified using the alias
+    ///      INSERT (X.ID, X.NAME)
+    ///      VALUES (FOO_IMP.ID, UPPER(FOO_IMP.NAME))
+    /// ```
+    ///
+    /// Note: in the latter case, the qualifier must match the target table
+    /// name or its alias if one is present. The parser will enforce this.
+    ///
+    /// The default implementation always returns `false` not allowing the
+    /// qualifiers.
+    fn supports_merge_insert_qualified_columns(&self) -> bool {
+        false
+    }
+
+    /// Returns `true` if the dialect supports specify an INSERT predicate in
+    /// MERGE statements. Example:
+    ///
+    /// ```sql
+    /// MERGE INTO FOO
+    /// USING FOO_IMP
+    ///    ON (FOO.ID = FOO_IMP.ID)
+    ///  WHEN NOT MATCHED THEN
+    ///      INSERT (ID, NAME)
+    ///      VALUES (FOO_IMP.ID, UPPER(FOO_IMP.NAME))
+    ///       -- insert predicate
+    ///       WHERE NOT FOO_IMP.NAME like '%.IGNORE'
+    /// ```
+    ///
+    /// The default implementation always returns `false` indicating no
+    /// support for the additional predicate.
+    ///
+    /// See also [Dialect::supports_merge_update_predicate] and
+    /// [Dialect::supports_merge_update_delete_predicate].
+    fn supports_merge_insert_predicate(&self) -> bool {
+        false
+    }
+
+    /// Indicates the supports of UPDATE predicates in MERGE
+    /// statements. Example:
+    ///
+    /// ```sql
+    /// MERGE INTO FOO
+    /// USING FOO_IMPORT
+    ///    ON (FOO.ID = FOO_IMPORT.ID)
+    ///  WHEN MATCHED THEN
+    ///      UPDATE SET FOO.NAME = FOO_IMPORT.NAME
+    ///       -- update predicate
+    ///       WHERE FOO.NAME <> 'pete'
+    /// ```
+    ///
+    /// The default implementation always returns false indicating no support
+    /// for the additional predicate.
+    ///
+    /// See also [Dialect::supports_merge_insert_predicate] and
+    /// [Dialect::supports_merge_update_delete_predicate].
+    fn supports_merge_update_predicate(&self) -> bool {
+        false
+    }
+
+    /// Indicates the supports of UPDATE ... DELETEs and associated predicates
+    /// in MERGE statements. Example:
+    ///
+    /// ```sql
+    /// MERGE INTO FOO
+    /// USING FOO_IMPORT
+    ///    ON (FOO.ID = FOO_IMPORT.ID)
+    ///  WHEN MATCHED THEN
+    ///      UPDATE SET FOO.NAME = FOO_IMPORT.NAME
+    ///      -- update delete with predicate
+    ///      DELETE WHERE UPPER(FOO.NAME) == FOO.NAME
+    /// ```
+    ///
+    /// The default implementation always returns false indicating no support
+    /// for the `UPDATE ... DELETE` and its associated predicate.
+    ///
+    /// See also [Dialect::supports_merge_insert_predicate] and
+    /// [Dialect::supports_merge_update_predicate].
+    fn supports_merge_update_delete_predicate(&self) -> bool {
+        false
+    }
+
     /// Dialect-specific infix parser override
     ///
     /// This method is called to parse the next infix expression.