Merge remote-tracking branch 'origin/feat/exclude-constraint-upstream' into feat/exclude-constraint-upstream

Author: fmguerreiro

src/ast/table_constraints.rs

@@ -20,13 +20,13 @@
 use crate::ast::{
     display_comma_separated, display_separated, ConstraintCharacteristics,
     ConstraintReferenceMatchKind, Expr, Ident, IndexColumn, IndexOption, IndexType,
-    KeyOrIndexDisplay, NullsDistinctOption, ObjectName, ReferentialAction,
+    KeyOrIndexDisplay, NullsDistinctOption, ObjectName, OrderByOptions, ReferentialAction,
 };
 use crate::tokenizer::Span;
 use core::fmt;
 
 #[cfg(not(feature = "std"))]
-use alloc::{boxed::Box, vec::Vec};
+use alloc::{boxed::Box, string::String, vec::Vec};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
@@ -619,7 +619,7 @@ impl crate::ast::Spanned for ConstraintUsingIndex {
 
 /// One element in an `EXCLUDE` constraint's element list.
 ///
-/// `<expr> WITH <operator>`
+/// `{ column_name | ( expression ) } [ opclass ] [ ASC | DESC ] [ NULLS { FIRST | LAST } ] WITH <operator>`
 ///
 /// See <https://www.postgresql.org/docs/current/sql-createtable.html#SQL-CREATETABLE-EXCLUDE>
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
@@ -628,20 +628,32 @@ impl crate::ast::Spanned for ConstraintUsingIndex {
 pub struct ExclusionElement {
     /// The index expression or column name.
     pub expr: Expr,
-    /// The exclusion operator (e.g. `&&`, `<->`, `=`).
+    /// Optional operator class (e.g. `gist_geometry_ops_nd`).
+    pub operator_class: Option<ObjectName>,
+    /// Ordering options (ASC/DESC, NULLS FIRST/LAST).
+    pub order: OrderByOptions,
+    /// The exclusion operator. Either a simple token (`&&`, `=`, `<->`) or the
+    /// Postgres schema-qualified form `OPERATOR(schema.op)`.
     pub operator: String,
 }
 
 impl fmt::Display for ExclusionElement {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{} WITH {}", self.expr, self.operator)
+        write!(f, "{}", self.expr)?;
+        if let Some(opclass) = &self.operator_class {
+            write!(f, " {opclass}")?;
+        }
+        write!(f, "{} WITH {}", self.order, self.operator)
     }
 }
 
 impl crate::ast::Spanned for ExclusionElement {
     fn span(&self) -> Span {
-        // Operator is stored as a plain String with no source span; only expr contributes.
-        self.expr.span()
+        let mut span = self.expr.span();
+        if let Some(opclass) = &self.operator_class {
+            span = span.union(&opclass.span());
+        }
+        span
     }
 }
 
@@ -691,11 +703,7 @@ impl fmt::Display for ExclusionConstraint {
 
 impl crate::ast::Spanned for ExclusionConstraint {
     fn span(&self) -> Span {
-        fn union_spans<I: Iterator<Item = Span>>(iter: I) -> Span {
-            Span::union_iter(iter)
-        }
-
-        union_spans(
+        Span::union_iter(
             self.name
                 .iter()
                 .map(|i| i.span)

src/parser/mod.rs

@@ -9905,16 +9905,18 @@ impl<'a> Parser<'a> {
                     .into(),
                 ))
             }
-            Token::Word(w) if w.keyword == Keyword::EXCLUDE => {
+            Token::Word(w)
+                if w.keyword == Keyword::EXCLUDE
+                    && dialect_of!(self is PostgreSqlDialect | GenericDialect) =>
+            {
                 let index_method = if self.parse_keyword(Keyword::USING) {
                     Some(self.parse_identifier()?)
                 } else {
                     None
                 };
 
                 self.expect_token(&Token::LParen)?;
-                let elements =
-                    self.parse_comma_separated(|p| p.parse_exclusion_element())?;
+                let elements = self.parse_comma_separated(|p| p.parse_exclusion_element())?;
                 self.expect_token(&Token::RParen)?;
 
                 let include = if self.parse_keyword(Keyword::INCLUDE) {
@@ -9961,20 +9963,54 @@ impl<'a> Parser<'a> {
     }
 
     fn parse_exclusion_element(&mut self) -> Result<ExclusionElement, ParserError> {
-        let expr = self.parse_expr()?;
+        // `index_elem` grammar: { col | (expr) } [ opclass ] [ ASC | DESC ] [ NULLS FIRST | LAST ].
+        // Shared with `CREATE INDEX` columns.
+        let (
+            OrderByExpr {
+                expr,
+                options: order,
+                ..
+            },
+            operator_class,
+        ) = self.parse_order_by_expr_inner(true)?;
+
         self.expect_keyword_is(Keyword::WITH)?;
+        let operator = self.parse_exclusion_operator()?;
+
+        Ok(ExclusionElement {
+            expr,
+            operator_class,
+            order,
+            operator,
+        })
+    }
+
+    /// Parse the operator that follows `WITH` in an `EXCLUDE` element.
+    ///
+    /// Accepts either a single operator token (e.g. `=`, `&&`, `<->`) or the
+    /// Postgres `OPERATOR(schema.op)` form for schema-qualified operators.
+    fn parse_exclusion_operator(&mut self) -> Result<String, ParserError> {
+        if self.parse_keyword(Keyword::OPERATOR) {
+            self.expect_token(&Token::LParen)?;
+            let mut parts = vec![];
+            loop {
+                self.advance_token();
+                parts.push(self.get_current_token().to_string());
+                if !self.consume_token(&Token::Period) {
+                    break;
+                }
+            }
+            self.expect_token(&Token::RParen)?;
+            return Ok(format!("OPERATOR({})", parts.join(".")));
+        }
+
         let operator_token = self.next_token();
         match &operator_token.token {
-            Token::EOF
-            | Token::RParen
-            | Token::Comma
-            | Token::SemiColon => {
-                return self.expected("exclusion operator", operator_token);
+            Token::EOF | Token::RParen | Token::Comma | Token::SemiColon => {
+                self.expected("exclusion operator", operator_token)
             }
-            _ => {}
+            _ => Ok(operator_token.token.to_string()),
         }
-        let operator = operator_token.token.to_string();
-        Ok(ExclusionElement { expr, operator })
     }
 
     fn parse_optional_nulls_distinct(&mut self) -> Result<NullsDistinctOption, ParserError> {