Snowflake: Add support for CONNECT_BY_ROOT (apache#1780)

Author: tomershaniii

src/ast/mod.rs

@@ -930,12 +930,14 @@ pub enum Expr {
     Nested(Box<Expr>),
     /// A literal value, such as string, number, date or NULL
     Value(ValueWithSpan),
+    /// Prefixed expression, e.g. introducer strings, projection prefix
     /// <https://dev.mysql.com/doc/refman/8.0/en/charset-introducer.html>
-    IntroducedString {
-        introducer: String,
+    /// <https://docs.snowflake.com/en/sql-reference/constructs/connect-by>
+    Prefixed {
+        prefix: Ident,
         /// The value of the constant.
         /// Hint: you can unwrap the string value using `value.into_string()`.
-        value: Value,
+        value: Box<Expr>,
     },
     /// A constant of form `<data_type> 'value'`.
     /// This can represent ANSI SQL `DATE`, `TIME`, and `TIMESTAMP` literals (such as `DATE '2020-01-01'`),
@@ -1655,7 +1657,7 @@ impl fmt::Display for Expr {
             Expr::Collate { expr, collation } => write!(f, "{expr} COLLATE {collation}"),
             Expr::Nested(ast) => write!(f, "({ast})"),
             Expr::Value(v) => write!(f, "{v}"),
-            Expr::IntroducedString { introducer, value } => write!(f, "{introducer} {value}"),
+            Expr::Prefixed { prefix, value } => write!(f, "{prefix} {value}"),
             Expr::TypedString { data_type, value } => {
                 write!(f, "{data_type}")?;
                 write!(f, " {value}")

src/ast/spans.rs

@@ -1543,7 +1543,7 @@ impl Spanned for Expr {
                             .map(|items| union_spans(items.iter().map(|i| i.span()))),
                     ),
             ),
-            Expr::IntroducedString { value, .. } => value.span(),
+            Expr::Prefixed { value, .. } => value.span(),
             Expr::Case {
                 operand,
                 conditions,

src/dialect/mod.rs

@@ -888,6 +888,12 @@ pub trait Dialect: Debug + Any {
         keywords::RESERVED_FOR_TABLE_FACTOR
     }
 
+    /// Returns reserved keywords that may prefix a select item expression
+    /// e.g. `SELECT CONNECT_BY_ROOT name FROM Tbl2` (Snowflake)
+    fn get_reserved_keywords_for_select_item_operator(&self) -> &[Keyword] {
+        &[]
+    }
+
     /// Returns true if this dialect supports the `TABLESAMPLE` option
     /// before the table alias option. For example:
     ///

src/dialect/snowflake.rs

@@ -44,6 +44,7 @@ use alloc::{format, vec};
 use super::keywords::RESERVED_FOR_IDENTIFIER;
 use sqlparser::ast::StorageSerializationPolicy;
 
+const RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR: [Keyword; 1] = [Keyword::CONNECT_BY_ROOT];
 /// A [`Dialect`] for [Snowflake](https://www.snowflake.com/)
 #[derive(Debug, Default)]
 pub struct SnowflakeDialect;
@@ -346,6 +347,11 @@ impl Dialect for SnowflakeDialect {
     fn supports_group_by_expr(&self) -> bool {
         true
     }
+
+    /// See: <https://docs.snowflake.com/en/sql-reference/constructs/connect-by>
+    fn get_reserved_keywords_for_select_item_operator(&self) -> &[Keyword] {
+        &RESERVED_KEYWORDS_FOR_SELECT_ITEM_OPERATOR
+    }
 }
 
 fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {

src/keywords.rs

@@ -207,6 +207,7 @@ define_keywords!(
     CONNECT,
     CONNECTION,
     CONNECTOR,
+    CONNECT_BY_ROOT,
     CONSTRAINT,
     CONTAINS,
     CONTINUE,

src/parser/mod.rs

@@ -1388,9 +1388,9 @@ impl<'a> Parser<'a> {
             | Token::HexStringLiteral(_)
                 if w.value.starts_with('_') =>
             {
-                Ok(Expr::IntroducedString {
-                    introducer: w.value.clone(),
-                    value: self.parse_introduced_string_value()?,
+                Ok(Expr::Prefixed {
+                    prefix: w.clone().into_ident(w_span),
+                    value: self.parse_introduced_string_expr()?.into(),
                 })
             }
             // string introducer https://dev.mysql.com/doc/refman/8.0/en/charset-introducer.html
@@ -1399,9 +1399,9 @@ impl<'a> Parser<'a> {
             | Token::HexStringLiteral(_)
                 if w.value.starts_with('_') =>
             {
-                Ok(Expr::IntroducedString {
-                    introducer: w.value.clone(),
-                    value: self.parse_introduced_string_value()?,
+                Ok(Expr::Prefixed {
+                    prefix: w.clone().into_ident(w_span),
+                    value: self.parse_introduced_string_expr()?.into(),
                 })
             }
             Token::Arrow if self.dialect.supports_lambda_functions() => {
@@ -9035,13 +9035,19 @@ impl<'a> Parser<'a> {
         }
     }
 
-    fn parse_introduced_string_value(&mut self) -> Result<Value, ParserError> {
+    fn parse_introduced_string_expr(&mut self) -> Result<Expr, ParserError> {
         let next_token = self.next_token();
         let span = next_token.span;
         match next_token.token {
-            Token::SingleQuotedString(ref s) => Ok(Value::SingleQuotedString(s.to_string())),
-            Token::DoubleQuotedString(ref s) => Ok(Value::DoubleQuotedString(s.to_string())),
-            Token::HexStringLiteral(ref s) => Ok(Value::HexStringLiteral(s.to_string())),
+            Token::SingleQuotedString(ref s) => Ok(Expr::Value(
+                Value::SingleQuotedString(s.to_string()).with_span(span),
+            )),
+            Token::DoubleQuotedString(ref s) => Ok(Expr::Value(
+                Value::DoubleQuotedString(s.to_string()).with_span(span),
+            )),
+            Token::HexStringLiteral(ref s) => Ok(Expr::Value(
+                Value::HexStringLiteral(s.to_string()).with_span(span),
+            )),
             unexpected => self.expected(
                 "a string value",
                 TokenWithSpan {
@@ -13968,6 +13974,13 @@ impl<'a> Parser<'a> {
 
     /// Parse a comma-delimited list of projections after SELECT
     pub fn parse_select_item(&mut self) -> Result<SelectItem, ParserError> {
+        let prefix = self
+            .parse_one_of_keywords(
+                self.dialect
+                    .get_reserved_keywords_for_select_item_operator(),
+            )
+            .map(|keyword| Ident::new(format!("{:?}", keyword)));
+
         match self.parse_wildcard_expr()? {
             Expr::QualifiedWildcard(prefix, token) => Ok(SelectItem::QualifiedWildcard(
                 SelectItemQualifiedWildcardKind::ObjectName(prefix),
@@ -14012,8 +14025,11 @@ impl<'a> Parser<'a> {
             expr => self
                 .maybe_parse_select_item_alias()
                 .map(|alias| match alias {
-                    Some(alias) => SelectItem::ExprWithAlias { expr, alias },
-                    None => SelectItem::UnnamedExpr(expr),
+                    Some(alias) => SelectItem::ExprWithAlias {
+                        expr: maybe_prefixed_expr(expr, prefix),
+                        alias,
+                    },
+                    None => SelectItem::UnnamedExpr(maybe_prefixed_expr(expr, prefix)),
                 }),
         }
     }
@@ -15375,6 +15391,17 @@ impl<'a> Parser<'a> {
     }
 }
 
+fn maybe_prefixed_expr(expr: Expr, prefix: Option<Ident>) -> Expr {
+    if let Some(prefix) = prefix {
+        Expr::Prefixed {
+            prefix,
+            value: Box::new(expr),
+        }
+    } else {
+        expr
+    }
+}
+
 impl Word {
     #[deprecated(since = "0.54.0", note = "please use `into_ident` instead")]
     pub fn to_ident(&self, span: Span) -> Ident {

tests/sqlparser_mysql.rs

@@ -3020,9 +3020,12 @@ fn parse_hex_string_introducer() {
                 distinct: None,
                 top: None,
                 top_before_distinct: false,
-                projection: vec![SelectItem::UnnamedExpr(Expr::IntroducedString {
-                    introducer: "_latin1".to_string(),
-                    value: Value::HexStringLiteral("4D7953514C".to_string())
+                projection: vec![SelectItem::UnnamedExpr(Expr::Prefixed {
+                    prefix: Ident::from("_latin1"),
+                    value: Expr::Value(
+                        Value::HexStringLiteral("4D7953514C".to_string()).with_empty_span()
+                    )
+                    .into(),
                 })],
                 from: vec![],
                 lateral_views: vec![],

tests/sqlparser_snowflake.rs

@@ -3983,3 +3983,45 @@ fn test_nested_join_without_parentheses() {
         }],
     );
 }
+
+#[test]
+fn parse_connect_by_root_operator() {
+    let sql = "SELECT CONNECT_BY_ROOT name AS root_name FROM Tbl1";
+
+    match snowflake().verified_stmt(sql) {
+        Statement::Query(query) => {
+            assert_eq!(
+                query.body.as_select().unwrap().projection[0],
+                SelectItem::ExprWithAlias {
+                    expr: Expr::Prefixed {
+                        prefix: Ident::new("CONNECT_BY_ROOT"),
+                        value: Box::new(Expr::Identifier(Ident::new("name")))
+                    },
+                    alias: Ident::new("root_name"),
+                }
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    let sql = "SELECT CONNECT_BY_ROOT name FROM Tbl2";
+    match snowflake().verified_stmt(sql) {
+        Statement::Query(query) => {
+            assert_eq!(
+                query.body.as_select().unwrap().projection[0],
+                SelectItem::UnnamedExpr(Expr::Prefixed {
+                    prefix: Ident::new("CONNECT_BY_ROOT"),
+                    value: Box::new(Expr::Identifier(Ident::new("name")))
+                })
+            );
+        }
+        _ => unreachable!(),
+    }
+
+    let sql = "SELECT CONNECT_BY_ROOT FROM Tbl2";
+    let res = snowflake().parse_sql_statements(sql);
+    assert_eq!(
+        res.unwrap_err().to_string(),
+        "sql parser error: Expected an expression, found: FROM"
+    );
+}