snowflake: handle pivot

Author: donhcd

src/ast/query.rs

@@ -267,6 +267,12 @@ pub enum TableFactor {
         expr: Expr,
         alias: Option<TableAlias>,
     },
+    /// https://docs.snowflake.com/en/sql-reference/constructs/pivot.html
+    Pivot {
+        expr: Expr,
+        val: Ident,
+        pivot_vals: Vec<Expr>,
+    },
     /// Represents a parenthesized table factor. The SQL spec only allows a
     /// join expression (`(foo <JOIN> bar [ <JOIN> baz ... ])`) to be nested,
     /// possibly several times.
@@ -332,6 +338,19 @@ impl fmt::Display for TableFactor {
                 }
                 Ok(())
             }
+            TableFactor::Pivot {
+                expr,
+                val,
+                pivot_vals,
+            } => {
+                write!(f, "({} FOR {} IN (", expr, val)?;
+                let mut delim = "";
+                for pivot_val in pivot_vals {
+                    write!(f, "{}{}", delim, pivot_val)?;
+                    delim = ", ";
+                }
+                write!(f, "))")
+            }
             TableFactor::NestedJoin(table_reference) => write!(f, "({})", table_reference),
         }
     }
@@ -414,6 +433,8 @@ impl fmt::Display for Join {
                 self.relation,
                 suffix(constraint)
             ),
+            JoinOperator::Pivot => write!(f, " PIVOT {}", self.relation),
+            JoinOperator::Unpivot => write!(f, " UNPIVOT {}", self.relation),
             JoinOperator::CrossJoin => write!(f, " CROSS JOIN {}", self.relation),
             JoinOperator::CrossApply => write!(f, " CROSS APPLY {}", self.relation),
             JoinOperator::OuterApply => write!(f, " OUTER APPLY {}", self.relation),
@@ -428,6 +449,11 @@ pub enum JoinOperator {
     LeftOuter(JoinConstraint),
     RightOuter(JoinConstraint),
     FullOuter(JoinConstraint),
+    /// [UN]PIVOT not actually a join but it seems to fit here syntactically
+    /// https://docs.snowflake.com/en/sql-reference/constructs/pivot.html
+    Pivot,
+    /// https://docs.snowflake.com/en/sql-reference/constructs/unpivot.html
+    Unpivot,
     CrossJoin,
     /// CROSS APPLY (non-standard)
     CrossApply,

src/dialect/keywords.rs

@@ -322,6 +322,7 @@ define_keywords!(
     PERCENTILE_DISC,
     PERCENT_RANK,
     PERIOD,
+    PIVOT,
     PORTION,
     POSITION,
     POSITION_REGEX,
@@ -441,6 +442,7 @@ define_keywords!(
     UNIQUE,
     UNKNOWN,
     UNNEST,
+    UNPIVOT,
     UPDATE,
     UPPER,
     USER,

src/parser.rs

@@ -1802,10 +1802,17 @@ impl<'a> Parser<'a> {
         &mut self,
         reserved_kwds: &[Keyword],
     ) -> Result<Option<TableAlias>, ParserError> {
+        let orig_index = self.index;
         match self.parse_optional_alias(reserved_kwds)? {
             Some(name) => {
-                let columns = self.parse_parenthesized_column_list(Optional)?;
-                Ok(Some(TableAlias { name, columns }))
+                if let Ok(columns) = self.parse_parenthesized_column_list(Optional) {
+                    Ok(Some(TableAlias { name, columns }))
+                } else {
+                    // if column list doesn't parse correctly, reset back to
+                    // original state
+                    self.index = orig_index;
+                    Ok(None)
+                }
             }
             None => Ok(None),
         }
@@ -2176,6 +2183,16 @@ impl<'a> Parser<'a> {
                     relation: self.parse_table_factor()?,
                     join_operator: JoinOperator::OuterApply,
                 }
+            } else if self.parse_keyword(Keyword::PIVOT) {
+                Join {
+                    relation: self.parse_pivot_body()?,
+                    join_operator: JoinOperator::Pivot,
+                }
+            } else if self.parse_keyword(Keyword::UNPIVOT) {
+                Join {
+                    relation: self.parse_pivot_body()?,
+                    join_operator: JoinOperator::Unpivot,
+                }
             } else {
                 let natural = self.parse_keyword(Keyword::NATURAL);
                 let peek_keyword = if let Token::Word(w) = self.peek_token() {
@@ -2315,6 +2332,7 @@ impl<'a> Parser<'a> {
                             // `(mytable AS alias)`
                             alias.replace(outer_alias);
                         }
+                        TableFactor::Pivot { .. } => unreachable!(),
                         TableFactor::NestedJoin(_) => unreachable!(),
                     };
                 }
@@ -2354,6 +2372,23 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_pivot_body(&mut self) -> Result<TableFactor, ParserError> {
+        self.expect_token(&Token::LParen)?;
+        let expr = self.parse_expr()?;
+        self.expect_keyword(Keyword::FOR)?;
+        let val = self.parse_identifier()?;
+        self.expect_keyword(Keyword::IN)?;
+        self.expect_token(&Token::LParen)?;
+        let pivot_vals = self.parse_comma_separated(Parser::parse_expr)?;
+        self.expect_token(&Token::RParen)?;
+        self.expect_token(&Token::RParen)?;
+        Ok(TableFactor::Pivot {
+            expr,
+            val,
+            pivot_vals,
+        })
+    }
+
     pub fn parse_derived_table_factor(
         &mut self,
         lateral: IsLateral,