snowflake: support for CONNECT BY and data as a json access key (apache#18)

A customer needs these fixes now, so need to merge these directly into
our fork. When the upstream PRs land we'll revert this and incorporate
that version in the next release merge.

Author: jmhain

src/ast/mod.rs

@@ -39,8 +39,8 @@ pub use self::ddl::{
 };
 pub use self::operator::{BinaryOperator, UnaryOperator};
 pub use self::query::{
-    Cte, CteAsMaterialized, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch, ForClause,
-    ForJson, ForXml, GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinOperator,
+    ConnectBy, Cte, CteAsMaterialized, Distinct, ExceptSelectItem, ExcludeSelectItem, Fetch,
+    ForClause, ForJson, ForXml, GroupByExpr, IdentWithAlias, Join, JoinConstraint, JoinOperator,
     JsonTableColumn, JsonTableColumnErrorHandling, LateralView, LockClause, LockType,
     NamedWindowDefinition, NonBlock, Offset, OffsetRows, OrderByExpr, Query, RenameSelectItem,
     ReplaceSelectElement, ReplaceSelectItem, Select, SelectInto, SelectItem, SetExpr, SetOperator,
@@ -735,6 +735,8 @@ pub enum Expr {
     ///
     /// See <https://docs.snowflake.com/en/sql-reference/constructs/where#joins-in-the-where-clause>.
     OuterJoin(Box<Expr>),
+    /// A reference to the prior level in a CONNECT BY clause.
+    Prior(Box<Expr>),
 }
 
 impl fmt::Display for CastFormat {
@@ -1210,6 +1212,7 @@ impl fmt::Display for Expr {
             Expr::OuterJoin(expr) => {
                 write!(f, "{expr} (+)")
             }
+            Expr::Prior(expr) => write!(f, "PRIOR {expr}"),
         }
     }
 }

src/ast/query.rs

@@ -92,6 +92,8 @@ impl fmt::Display for Query {
 pub enum SetExpr {
     /// Restricted SELECT .. FROM .. HAVING (no ORDER BY or set operations)
     Select(Box<Select>),
+    /// SELECT .. FROM .. STARTING WITH .. CONNECT BY
+    ConnectBy(ConnectBy),
     /// Parenthesized SELECT subquery, which may include more set operations
     /// in its body and an optional ORDER BY / LIMIT.
     Query(Box<Query>),
@@ -112,6 +114,7 @@ impl fmt::Display for SetExpr {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             SetExpr::Select(s) => write!(f, "{s}"),
+            SetExpr::ConnectBy(c) => write!(f, "{c}"),
             SetExpr::Query(q) => write!(f, "({q})"),
             SetExpr::Values(v) => write!(f, "{v}"),
             SetExpr::Insert(v) => write!(f, "{v}"),
@@ -702,6 +705,36 @@ impl fmt::Display for TableWithJoins {
     }
 }
 
+/// Joins a table to itself to process hierarchical data in the table.
+///
+/// See <https://docs.snowflake.com/en/sql-reference/constructs/connect-by>.
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct ConnectBy {
+    /// SELECT
+    pub projection: Vec<SelectItem>,
+    /// FROM
+    pub from: Vec<TableWithJoins>,
+    /// START WITH
+    pub condition: Expr,
+    /// CONNECT BY
+    pub relationships: Vec<Expr>,
+}
+
+impl fmt::Display for ConnectBy {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(
+            f,
+            "SELECT {projection} FROM {from} START WITH {condition} CONNECT BY {relationships}",
+            projection = display_comma_separated(&self.projection),
+            from = display_comma_separated(&self.from),
+            condition = self.condition,
+            relationships = display_comma_separated(&self.relationships)
+        )
+    }
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

src/keywords.rs

@@ -797,6 +797,9 @@ pub const RESERVED_FOR_TABLE_ALIAS: &[Keyword] = &[
     Keyword::FOR,
     // for MYSQL PARTITION SELECTION
     Keyword::PARTITION,
+    // for Snowflake START WITH .. CONNECT BY
+    Keyword::START,
+    Keyword::CONNECT,
 ];
 
 /// Can't be used as a column alias, so that `SELECT <expr> alias`

src/parser/mod.rs

@@ -256,10 +256,22 @@ impl ParserOptions {
     }
 }
 
+#[derive(Copy, Clone)]
+enum ParserState {
+    /// The default state of the parser.
+    Normal,
+    /// The state when parsing a CONNECT BY expression. This allows parsing
+    /// PRIOR expressions while still allowing prior as an identifier name
+    /// in other contexts.
+    ConnectBy,
+}
+
 pub struct Parser<'a> {
     tokens: Vec<TokenWithLocation>,
     /// The index of the first unprocessed token in `self.tokens`
     index: usize,
+    /// The current state of the parser.
+    state: ParserState,
     /// The current dialect to use
     dialect: &'a dyn Dialect,
     /// Additional options that allow you to mix & match behavior
@@ -290,6 +302,7 @@ impl<'a> Parser<'a> {
         Self {
             tokens: vec![],
             index: 0,
+            state: ParserState::Normal,
             dialect,
             recursion_counter: RecursionCounter::new(DEFAULT_REMAINING_DEPTH),
             options: ParserOptions::default(),
@@ -966,6 +979,10 @@ impl<'a> Parser<'a> {
                     self.prev_token();
                     self.parse_bigquery_struct_literal()
                 }
+                Keyword::PRIOR if matches!(self.state, ParserState::ConnectBy) => {
+                    let expr = self.parse_subexpr(Self::PLUS_MINUS_PREC)?;
+                    Ok(Expr::Prior(Box::new(expr)))
+                }
                 // Here `w` is a word, check if it's a part of a multi-part
                 // identifier, a function call, or a simple identifier:
                 _ => match self.peek_token().token {
@@ -5894,12 +5911,15 @@ impl<'a> Parser<'a> {
                     )?,
                 },
                 // Case when Snowflake Semi-structured data like key:value
+                // FIXME: Find a better way to deal with this than exhaustively
+                //        listing every keyword allowed here.
                 Keyword::NoKeyword
                 | Keyword::LOCATION
                 | Keyword::TYPE
                 | Keyword::DATE
                 | Keyword::START
                 | Keyword::END
+                | Keyword::DATA
                     if dialect_of!(self is SnowflakeDialect | GenericDialect) =>
                 {
                     Ok(Value::UnQuotedString(w.value))
@@ -7154,7 +7174,7 @@ impl<'a> Parser<'a> {
         // We parse the expression using a Pratt parser, as in `parse_expr()`.
         // Start by parsing a restricted SELECT or a `(subquery)`:
         let mut expr = if self.parse_keyword(Keyword::SELECT) {
-            SetExpr::Select(Box::new(self.parse_select()?))
+            self.parse_select()?
         } else if self.consume_token(&Token::LParen) {
             // CTEs are not allowed here, but the parser currently accepts them
             let subquery = self.parse_query()?;
@@ -7242,7 +7262,7 @@ impl<'a> Parser<'a> {
 
     /// Parse a restricted `SELECT` statement (no CTEs / `UNION` / `ORDER BY`),
     /// assuming the initial `SELECT` was already consumed
-    pub fn parse_select(&mut self) -> Result<Select, ParserError> {
+    pub fn parse_select(&mut self) -> Result<SetExpr, ParserError> {
         let value_table_mode =
             if dialect_of!(self is BigQueryDialect) && self.parse_keyword(Keyword::AS) {
                 if self.parse_keyword(Keyword::VALUE) {
@@ -7294,6 +7314,18 @@ impl<'a> Parser<'a> {
             vec![]
         };
 
+        if distinct.is_none()
+            && top.is_none()
+            && into.is_none()
+            && !from.is_empty()
+            && self
+                .parse_one_of_keywords(&[Keyword::START, Keyword::CONNECT])
+                .is_some()
+        {
+            self.prev_token();
+            return Ok(SetExpr::ConnectBy(self.parse_connect_by(projection, from)?));
+        }
+
         let mut lateral_views = vec![];
         loop {
             if self.parse_keywords(&[Keyword::LATERAL, Keyword::VIEW]) {
@@ -7377,7 +7409,7 @@ impl<'a> Parser<'a> {
             None
         };
 
-        Ok(Select {
+        Ok(SetExpr::Select(Box::new(Select {
             distinct,
             top,
             projection,
@@ -7393,6 +7425,48 @@ impl<'a> Parser<'a> {
             named_window: named_windows,
             qualify,
             value_table_mode,
+        })))
+    }
+
+    fn with_state<T, F>(&mut self, state: ParserState, mut f: F) -> Result<T, ParserError>
+    where
+        F: FnMut(&mut Parser) -> Result<T, ParserError>,
+    {
+        let current_state = self.state;
+        self.state = state;
+        let res = f(self);
+        self.state = current_state;
+        res
+    }
+
+    pub fn parse_connect_by(
+        &mut self,
+        projection: Vec<SelectItem>,
+        from: Vec<TableWithJoins>,
+    ) -> Result<ConnectBy, ParserError> {
+        debug_assert!(!from.is_empty());
+
+        let (condition, relationships) = if self.parse_keywords(&[Keyword::CONNECT, Keyword::BY]) {
+            let relationships = self.with_state(ParserState::ConnectBy, |parser| {
+                parser.parse_comma_separated(Parser::parse_expr)
+            })?;
+            self.expect_keywords(&[Keyword::START, Keyword::WITH])?;
+            let condition = self.parse_expr()?;
+            (condition, relationships)
+        } else {
+            self.expect_keywords(&[Keyword::START, Keyword::WITH])?;
+            let condition = self.parse_expr()?;
+            self.expect_keywords(&[Keyword::CONNECT, Keyword::BY])?;
+            let relationships = self.with_state(ParserState::ConnectBy, |parser| {
+                parser.parse_comma_separated(Parser::parse_expr)
+            })?;
+            (condition, relationships)
+        };
+        Ok(ConnectBy {
+            projection,
+            from,
+            condition,
+            relationships,
         })
     }
 

src/test_utils.rs

@@ -157,6 +157,16 @@ impl TestedDialects {
         }
     }
 
+    /// Ensures that `sql` parses as a single [Query], and that
+    /// re-serializing the parse result matches the given canonical
+    /// sql string.
+    pub fn verified_query_with_canonical(&self, query: &str, canonical: &str) -> Query {
+        match self.one_statement_parses_to(query, canonical) {
+            Statement::Query(query) => *query,
+            _ => panic!("Expected Query"),
+        }
+    }
+
     /// Ensures that `sql` parses as a single [Select], and that
     /// re-serializing the parse result produces the same `sql`
     /// string (is not modified after a serialization round-trip).

tests/sqlparser_snowflake.rs

@@ -230,7 +230,7 @@ fn parse_json_using_colon() {
 
     snowflake().one_statement_parses_to("SELECT a:b::int FROM t", "SELECT CAST(a:b AS INT) FROM t");
 
-    let sql = "SELECT a:start, a:end FROM t";
+    let sql = "SELECT a:start, a:end, a:data FROM t";
     let select = snowflake().verified_only_select(sql);
     assert_eq!(
         vec![
@@ -243,6 +243,11 @@ fn parse_json_using_colon() {
                 left: Box::new(Expr::Identifier(Ident::new("a"))),
                 operator: JsonOperator::Colon,
                 right: Box::new(Expr::Value(Value::UnQuotedString("end".to_string()))),
+            }),
+            SelectItem::UnnamedExpr(Expr::JsonAccess {
+                left: Box::new(Expr::Identifier(Ident::new("a"))),
+                operator: JsonOperator::Colon,
+                right: Box::new(Expr::Value(Value::UnQuotedString("data".to_string()))),
             })
         ],
         select.projection
@@ -1503,3 +1508,88 @@ fn parse_comma_outer_join() {
         "SELECT t1.c1, t2.c2 FROM t1, t2 WHERE t1.c1 = t2.c2 (+)",
     );
 }
+
+#[test]
+fn parse_connect_by() {
+    let expect_query = Query {
+        with: None,
+        body: Box::new(SetExpr::ConnectBy(ConnectBy {
+            projection: vec![
+                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("employee_id"))),
+                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("manager_id"))),
+                SelectItem::UnnamedExpr(Expr::Identifier(Ident::new("title"))),
+            ],
+            from: vec![TableWithJoins {
+                relation: TableFactor::Table {
+                    name: ObjectName(vec![Ident::new("employees")]),
+                    alias: None,
+                    args: None,
+                    with_hints: vec![],
+                    version: None,
+                    partitions: vec![],
+                },
+                joins: vec![],
+            }],
+            condition: Expr::BinaryOp {
+                left: Box::new(Expr::Identifier(Ident::new("title"))),
+                op: BinaryOperator::Eq,
+                right: Box::new(Expr::Value(Value::SingleQuotedString(
+                    "president".to_owned(),
+                ))),
+            },
+            relationships: vec![Expr::BinaryOp {
+                left: Box::new(Expr::Identifier(Ident::new("manager_id"))),
+                op: BinaryOperator::Eq,
+                right: Box::new(Expr::Prior(Box::new(Expr::Identifier(Ident::new(
+                    "employee_id",
+                ))))),
+            }],
+        })),
+        order_by: vec![OrderByExpr {
+            expr: Expr::Identifier(Ident::new("employee_id")),
+            asc: None,
+            nulls_first: None,
+        }],
+        limit: None,
+        limit_by: vec![],
+        offset: None,
+        fetch: None,
+        locks: vec![],
+        for_clause: None,
+    };
+
+    let connect_by_1 = concat!(
+        "SELECT employee_id, manager_id, title FROM employees ",
+        "START WITH title = 'president' ",
+        "CONNECT BY manager_id = PRIOR employee_id ",
+        "ORDER BY employee_id"
+    );
+
+    assert_eq!(
+        snowflake_and_generic().verified_query(connect_by_1),
+        expect_query
+    );
+
+    // CONNECT BY can come before START WITH
+    let connect_by_2 = concat!(
+        "SELECT employee_id, manager_id, title FROM employees ",
+        "CONNECT BY manager_id = PRIOR employee_id ",
+        "START WITH title = 'president' ",
+        "ORDER BY employee_id"
+    );
+    assert_eq!(
+        snowflake_and_generic().verified_query_with_canonical(connect_by_2, connect_by_1),
+        expect_query
+    );
+
+    // PRIOR expressions are only valid within a CONNECT BY, and the the token
+    // `prior` is valid as an identifier anywhere else.
+    assert_eq!(
+        snowflake_and_generic()
+            .verified_only_select("SELECT prior FROM some_table")
+            .projection,
+        vec![SelectItem::UnnamedExpr(Expr::Identifier(Ident::new(
+            "prior"
+        )))]
+    );
+}