Support parsing optional nulls handling for unique constraint (#1567)

apache · Dec 4, 2024 · 6517da6 · 6517da6
1 parent 6d4188d
commit 6517da6
Show file tree

Hide file tree

Showing 6 changed files with 71 additions and 4 deletions.
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
@@ -669,6 +669,8 @@ pub enum TableConstraint {
         columns: Vec<Ident>,
         index_options: Vec<IndexOption>,
         characteristics: Option<ConstraintCharacteristics>,
+        /// Optional Postgres nulls handling: `[ NULLS [ NOT ] DISTINCT ]`
+        nulls_distinct: NullsDistinctOption,
     },
     /// MySQL [definition][1] for `PRIMARY KEY` constraints statements:\
     /// * `[CONSTRAINT [<name>]] PRIMARY KEY [index_name] [index_type] (<columns>) <index_options>`
@@ -777,10 +779,11 @@ impl fmt::Display for TableConstraint {
                 columns,
                 index_options,
                 characteristics,
+                nulls_distinct,
             } => {
                 write!(
                     f,
-                    "{}UNIQUE{index_type_display:>}{}{} ({})",
+                    "{}UNIQUE{nulls_distinct}{index_type_display:>}{}{} ({})",
                     display_constraint_name(name),
                     display_option_spaced(index_name),
                     display_option(" USING ", "", index_type),
@@ -988,6 +991,31 @@ impl fmt::Display for IndexOption {
     }
 }
 
+/// [Postgres] unique index nulls handling option: `[ NULLS [ NOT ] DISTINCT ]`
+///
+/// [Postgres]: https://www.postgresql.org/docs/17/sql-altertable.html
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum NullsDistinctOption {
+    /// Not specified
+    None,
+    /// NULLS DISTINCT
+    Distinct,
+    /// NULLS NOT DISTINCT
+    NotDistinct,
+}
+
+impl fmt::Display for NullsDistinctOption {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Self::None => Ok(()),
+            Self::Distinct => write!(f, " NULLS DISTINCT"),
+            Self::NotDistinct => write!(f, " NULLS NOT DISTINCT"),
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -49,9 +49,10 @@ pub use self::ddl::{
     ClusteredBy, ColumnDef, ColumnOption, ColumnOptionDef, ColumnPolicy, ColumnPolicyProperty,
     ConstraintCharacteristics, CreateFunction, Deduplicate, DeferrableInitial, GeneratedAs,
     GeneratedExpressionMode, IdentityParameters, IdentityProperty, IdentityPropertyFormatKind,
-    IdentityPropertyKind, IdentityPropertyOrder, IndexOption, IndexType, KeyOrIndexDisplay, Owner,
-    Partition, ProcedureParam, ReferentialAction, TableConstraint, TagsColumnOption,
-    UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation, ViewColumnDef,
+    IdentityPropertyKind, IdentityPropertyOrder, IndexOption, IndexType, KeyOrIndexDisplay,
+    NullsDistinctOption, Owner, Partition, ProcedureParam, ReferentialAction, TableConstraint,
+    TagsColumnOption, UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation,
+    ViewColumnDef,
 };
 pub use self::dml::{CreateIndex, CreateTable, Delete, Insert};
 pub use self::operator::{BinaryOperator, UnaryOperator};

diff --git a/src/ast/spans.rs b/src/ast/spans.rs
@@ -587,6 +587,7 @@ impl Spanned for TableConstraint {
                 columns,
                 index_options: _,
                 characteristics,
+                nulls_distinct: _,
             } => union_spans(
                 name.iter()
                     .map(|i| i.span)

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -6729,6 +6729,8 @@ impl<'a> Parser<'a> {
                         .expected("`index_name` or `(column_name [, ...])`", self.peek_token());
                 }
 
+                let nulls_distinct = self.parse_optional_nulls_distinct()?;
+
                 // optional index name
                 let index_name = self.parse_optional_indent()?;
                 let index_type = self.parse_optional_using_then_index_type()?;
@@ -6744,6 +6746,7 @@ impl<'a> Parser<'a> {
                     columns,
                     index_options,
                     characteristics,
+                    nulls_distinct,
                 }))
             }
             Token::Word(w) if w.keyword == Keyword::PRIMARY => {
@@ -6866,6 +6869,20 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn parse_optional_nulls_distinct(&mut self) -> Result<NullsDistinctOption, ParserError> {
+        Ok(if self.parse_keyword(Keyword::NULLS) {
+            let not = self.parse_keyword(Keyword::NOT);
+            self.expect_keyword(Keyword::DISTINCT)?;
+            if not {
+                NullsDistinctOption::NotDistinct
+            } else {
+                NullsDistinctOption::Distinct
+            }
+        } else {
+            NullsDistinctOption::None
+        })
+    }
+
     pub fn maybe_parse_options(
         &mut self,
         keyword: Keyword,

diff --git a/tests/sqlparser_mysql.rs b/tests/sqlparser_mysql.rs
@@ -669,6 +669,7 @@ fn table_constraint_unique_primary_ctor(
             columns,
             index_options,
             characteristics,
+            nulls_distinct: NullsDistinctOption::None,
         },
         None => TableConstraint::PrimaryKey {
             name,

diff --git a/tests/sqlparser_postgres.rs b/tests/sqlparser_postgres.rs
@@ -594,6 +594,25 @@ fn parse_alter_table_constraints_rename() {
     }
 }
 
+#[test]
+fn parse_alter_table_constraints_unique_nulls_distinct() {
+    match pg_and_generic()
+        .verified_stmt("ALTER TABLE t ADD CONSTRAINT b UNIQUE NULLS NOT DISTINCT (c)")
+    {
+        Statement::AlterTable { operations, .. } => match &operations[0] {
+            AlterTableOperation::AddConstraint(TableConstraint::Unique {
+                nulls_distinct, ..
+            }) => {
+                assert_eq!(nulls_distinct, &NullsDistinctOption::NotDistinct)
+            }
+            _ => unreachable!(),
+        },
+        _ => unreachable!(),
+    }
+    pg_and_generic().verified_stmt("ALTER TABLE t ADD CONSTRAINT b UNIQUE NULLS DISTINCT (c)");
+    pg_and_generic().verified_stmt("ALTER TABLE t ADD CONSTRAINT b UNIQUE (c)");
+}
+
 #[test]
 fn parse_alter_table_disable() {
     pg_and_generic().verified_stmt("ALTER TABLE tab DISABLE ROW LEVEL SECURITY");