Snowflake: extend support for create table columns

validio-io · Sep 26, 2024 · d285272 · d285272
1 parent 8a534c0
commit d285272
Show file tree

Hide file tree

Showing 4 changed files with 114 additions and 24 deletions.
diff --git a/src/ast/ddl.rs b/src/ast/ddl.rs
@@ -1059,13 +1059,75 @@ impl fmt::Display for ColumnOptionDef {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub struct IdentityProperty {
+    pub format: IdentityPropertyCommand,
+    pub parameters: Option<IdentityParameters>,
+    pub order: Option<IdentityOrder>,
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum IdentityPropertyCommand {
+    Autoincrement,
+    Identity,
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct IdentityParameters {
+    pub format: IdentityParametersFormat,
     pub seed: Expr,
     pub increment: Expr,
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum IdentityParametersFormat {
+    FunctionCall,
+    StartIncrement,
+}
+
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum IdentityOrder {
+    Order,
+    Noorder,
+}
+
 impl fmt::Display for IdentityProperty {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "{}, {}", self.seed, self.increment)
+        if let Some(parameters) = &self.parameters {
+            write!(f, "{parameters}")?;
+        }
+        if let Some(order) = &self.order {
+            write!(f, "{order}")?;
+        }
+        Ok(())
+    }
+}
+
+impl fmt::Display for IdentityParameters {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self.format {
+            IdentityParametersFormat::FunctionCall => {
+                write!(f, "({}, {})", self.seed, self.increment)
+            }
+            IdentityParametersFormat::StartIncrement => {
+                write!(f, "START {} INCREMENT {})", self.seed, self.increment)
+            }
+        }
+    }
+}
+
+impl fmt::Display for IdentityOrder {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            IdentityOrder::Order => write!(f, "ORDER"),
+            IdentityOrder::Noorder => write!(f, "NOORDER"),
+        }
     }
 }
 
@@ -1139,13 +1201,14 @@ pub enum ColumnOption {
     /// [1]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#view_column_option_list
     /// [2]: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#column_option_list
     Options(Vec<SqlOption>),
-    /// MS SQL Server specific: Creates an identity column in a table.
+    /// Creates an identity or an autoincrement column in a table.
     /// Syntax
     /// ```sql
-    /// IDENTITY [ (seed , increment) ]
+    /// { IDENTITY | AUTOINCREMENT } [ (seed , increment) | START num INCREMENT num ] [ ORDER | NOORDER ]
     /// ```
     /// [MS SQL Server]: https://learn.microsoft.com/en-us/sql/t-sql/statements/create-table-transact-sql-identity-property
-    Identity(Option<IdentityProperty>),
+    /// [Snowflake]: https://docs.snowflake.com/en/sql-reference/sql/create-table
+    Identity(IdentityProperty),
 }
 
 impl fmt::Display for ColumnOption {
@@ -1248,10 +1311,11 @@ impl fmt::Display for ColumnOption {
                 write!(f, "OPTIONS({})", display_comma_separated(options))
             }
             Identity(parameters) => {
-                write!(f, "IDENTITY")?;
-                if let Some(parameters) = parameters {
-                    write!(f, "({parameters})")?;
+                match parameters.format {
+                    IdentityPropertyCommand::Identity => write!(f, "IDENTITY")?,
+                    IdentityPropertyCommand::Autoincrement => write!(f, "AUTOINCREMENT")?,
                 }
+                write!(f, "{parameters}")?;
                 Ok(())
             }
         }

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
@@ -41,8 +41,9 @@ pub use self::dcl::{AlterRoleOperation, ResetConfig, RoleOption, SetConfigValue,
 pub use self::ddl::{
     AlterColumnOperation, AlterIndexOperation, AlterTableOperation, ClusteredBy, ColumnDef,
     ColumnOption, ColumnOptionDef, ConstraintCharacteristics, Deduplicate, DeferrableInitial,
-    GeneratedAs, GeneratedExpressionMode, IdentityProperty, IndexOption, IndexType,
-    KeyOrIndexDisplay, Owner, Partition, ProcedureParam, ReferentialAction, TableConstraint,
+    GeneratedAs, GeneratedExpressionMode, IdentityParameters, IdentityParametersFormat,
+    IdentityProperty, IdentityPropertyCommand, IndexOption, IndexType, KeyOrIndexDisplay, Owner,
+    Partition, ProcedureParam, ReferentialAction, TableConstraint,
     UserDefinedTypeCompositeAttributeDef, UserDefinedTypeRepresentation, ViewColumnDef,
 };
 pub use self::dml::{CreateIndex, CreateTable, Delete, Insert};

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -6174,17 +6174,25 @@ impl<'a> Parser<'a> {
         } else if self.parse_keyword(Keyword::IDENTITY)
             && dialect_of!(self is MsSqlDialect | GenericDialect)
         {
-            let property = if self.consume_token(&Token::LParen) {
+            let parameters = if self.consume_token(&Token::LParen) {
                 let seed = self.parse_number()?;
                 self.expect_token(&Token::Comma)?;
                 let increment = self.parse_number()?;
                 self.expect_token(&Token::RParen)?;
 
-                Some(IdentityProperty { seed, increment })
+                Some(IdentityParameters {
+                    format: IdentityParametersFormat::FunctionCall,
+                    seed,
+                    increment,
+                })
             } else {
                 None
             };
-            Ok(Some(ColumnOption::Identity(property)))
+            Ok(Some(ColumnOption::Identity(IdentityProperty {
+                format: IdentityPropertyCommand::Identity,
+                parameters,
+                order: None,
+            })))
         } else {
             Ok(None)
         }

diff --git a/tests/sqlparser_mssql.rs b/tests/sqlparser_mssql.rs
@@ -921,7 +921,11 @@ fn parse_create_table_with_identity_column() {
             vec![
                 ColumnOptionDef {
                     name: None,
-                    option: ColumnOption::Identity(None),
+                    option: ColumnOption::Identity(IdentityProperty {
+                        format: IdentityPropertyCommand::Identity,
+                        parameters: None,
+                        order: None,
+                    }),
                 },
                 ColumnOptionDef {
                     name: None,
@@ -935,18 +939,31 @@ fn parse_create_table_with_identity_column() {
                 ColumnOptionDef {
                     name: None,
                     #[cfg(not(feature = "bigdecimal"))]
-                    option: ColumnOption::Identity(Some(IdentityProperty {
-                        seed: Expr::Value(Value::Number("1".to_string(), false)),
-                        increment: Expr::Value(Value::Number("1".to_string(), false)),
-                    })),
+                    option: ColumnOption::Identity(IdentityProperty {
+                        format: IdentityPropertyCommand::Identity,
+                        parameters: Some(IdentityParameters {
+                            format: IdentityParametersFormat::FunctionCall,
+                            seed: Expr::Value(Value::Number("1".to_string(), false)),
+                            increment: Expr::Value(Value::Number("1".to_string(), false)),
+                        }),
+                        order: None,
+                    }),
                     #[cfg(feature = "bigdecimal")]
-                    option: ColumnOption::Identity(Some(IdentityProperty {
-                        seed: Expr::Value(Value::Number(bigdecimal::BigDecimal::from(1), false)),
-                        increment: Expr::Value(Value::Number(
-                            bigdecimal::BigDecimal::from(1),
-                            false,
-                        )),
-                    })),
+                    option: ColumnOption::Identity(IdentityProperty {
+                        format: IdentityPropertyCommand::Identity,
+                        parameters: Some(IdentityParameters {
+                            format: IdentityParametersFormat::FunctionCall,
+                            seed: Expr::Value(Value::Number(
+                                bigdecimal::BigDecimal::from(1),
+                                false,
+                            )),
+                            increment: Expr::Value(Value::Number(
+                                bigdecimal::BigDecimal::from(1),
+                                false,
+                            )),
+                        }),
+                        order: None,
+                    }),
                 },
                 ColumnOptionDef {
                     name: None,