supporting snowflake extract syntax (#1374)

Co-authored-by: Andrew Lamb <[email protected]>

Author: seve-martinez

src/ast/mod.rs

@@ -477,6 +477,22 @@ pub enum CastKind {
     DoubleColon,
 }
 
+/// `EXTRACT` syntax variants.
+///
+/// In Snowflake dialect, the `EXTRACT` expression can support either the `from` syntax
+/// or the comma syntax.
+///
+/// See <https://docs.snowflake.com/en/sql-reference/functions/extract>
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum ExtractSyntax {
+    /// `EXTRACT( <date_or_time_part> FROM <date_or_time_expr> )`
+    From,
+    /// `EXTRACT( <date_or_time_part> , <date_or_timestamp_expr> )`
+    Comma,
+}
+
 /// An SQL expression of any type.
 ///
 /// The parser does not distinguish between expressions of different types
@@ -637,13 +653,15 @@ pub enum Expr {
         time_zone: Box<Expr>,
     },
     /// Extract a field from a timestamp e.g. `EXTRACT(MONTH FROM foo)`
+    /// Or `EXTRACT(MONTH, foo)`
     ///
     /// Syntax:
     /// ```sql
-    /// EXTRACT(DateTimeField FROM <expr>)
+    /// EXTRACT(DateTimeField FROM <expr>) | EXTRACT(DateTimeField, <expr>)
     /// ```
     Extract {
         field: DateTimeField,
+        syntax: ExtractSyntax,
         expr: Box<Expr>,
     },
     /// ```sql
@@ -1197,7 +1215,14 @@ impl fmt::Display for Expr {
                     write!(f, "{expr}::{data_type}")
                 }
             },
-            Expr::Extract { field, expr } => write!(f, "EXTRACT({field} FROM {expr})"),
+            Expr::Extract {
+                field,
+                syntax,
+                expr,
+            } => match syntax {
+                ExtractSyntax::From => write!(f, "EXTRACT({field} FROM {expr})"),
+                ExtractSyntax::Comma => write!(f, "EXTRACT({field}, {expr})"),
+            },
             Expr::Ceil { expr, field } => {
                 if field == &DateTimeField::NoDateTime {
                     write!(f, "CEIL({expr})")

src/parser/mod.rs

@@ -1682,12 +1682,25 @@ impl<'a> Parser<'a> {
     pub fn parse_extract_expr(&mut self) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
         let field = self.parse_date_time_field()?;
-        self.expect_keyword(Keyword::FROM)?;
+
+        let syntax = if self.parse_keyword(Keyword::FROM) {
+            ExtractSyntax::From
+        } else if self.consume_token(&Token::Comma)
+            && dialect_of!(self is SnowflakeDialect | GenericDialect)
+        {
+            ExtractSyntax::Comma
+        } else {
+            return Err(ParserError::ParserError(
+                "Expected 'FROM' or ','".to_string(),
+            ));
+        };
+
         let expr = self.parse_expr()?;
         self.expect_token(&Token::RParen)?;
         Ok(Expr::Extract {
             field,
             expr: Box::new(expr),
+            syntax,
         })
     }
 
@@ -1950,6 +1963,12 @@ impl<'a> Parser<'a> {
                 }
                 _ => self.expected("date/time field", next_token),
             },
+            Token::SingleQuotedString(_) if dialect_of!(self is SnowflakeDialect | GenericDialect) =>
+            {
+                self.prev_token();
+                let custom = self.parse_identifier(false)?;
+                Ok(DateTimeField::Custom(custom))
+            }
             _ => self.expected("date/time field", next_token),
         }
     }

tests/sqlparser_bigquery.rs

@@ -2136,6 +2136,7 @@ fn parse_extract_weekday() {
     assert_eq!(
         &Expr::Extract {
             field: DateTimeField::Week(Some(Ident::new("MONDAY"))),
+            syntax: ExtractSyntax::From,
             expr: Box::new(Expr::Identifier(Ident::new("d"))),
         },
         expr_from_projection(only(&select.projection)),

tests/sqlparser_common.rs

@@ -2430,6 +2430,7 @@ fn parse_extract() {
     assert_eq!(
         &Expr::Extract {
             field: DateTimeField::Year,
+            syntax: ExtractSyntax::From,
             expr: Box::new(Expr::Identifier(Ident::new("d"))),
         },
         expr_from_projection(only(&select.projection)),

tests/sqlparser_snowflake.rs

@@ -2019,6 +2019,35 @@ fn parse_extract_custom_part() {
     assert_eq!(
         &Expr::Extract {
             field: DateTimeField::Custom(Ident::new("eod")),
+            syntax: ExtractSyntax::From,
+            expr: Box::new(Expr::Identifier(Ident::new("d"))),
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+}
+
+#[test]
+fn parse_extract_comma() {
+    let sql = "SELECT EXTRACT(HOUR, d)";
+    let select = snowflake_and_generic().verified_only_select(sql);
+    assert_eq!(
+        &Expr::Extract {
+            field: DateTimeField::Hour,
+            syntax: ExtractSyntax::Comma,
+            expr: Box::new(Expr::Identifier(Ident::new("d"))),
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+}
+
+#[test]
+fn parse_extract_comma_quoted() {
+    let sql = "SELECT EXTRACT('hour', d)";
+    let select = snowflake_and_generic().verified_only_select(sql);
+    assert_eq!(
+        &Expr::Extract {
+            field: DateTimeField::Custom(Ident::with_quote('\'', "hour")),
+            syntax: ExtractSyntax::Comma,
             expr: Box::new(Expr::Identifier(Ident::new("d"))),
         },
         expr_from_projection(only(&select.projection)),