Add support for AT TIME ZONE (apache#539)

* Support for empty array literals

* Added support for AT TIME ZONE

Co-authored-by: Chris Allen <chrisa@indeed.com>

Author: bitemyapp

src/ast/mod.rs

@@ -285,6 +285,11 @@ pub enum Expr {
         expr: Box<Expr>,
         data_type: DataType,
     },
+    /// AT a timestamp to a different timezone e.g. `FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00'`
+    AtTimeZone {
+        timestamp: Box<Expr>,
+        time_zone: String,
+    },
     /// EXTRACT(DateTimeField FROM <expr>)
     Extract {
         field: DateTimeField,
@@ -562,6 +567,12 @@ impl fmt::Display for Expr {
             Expr::CompositeAccess { expr, key } => {
                 write!(f, "{}.{}", expr, key)
             }
+            Expr::AtTimeZone {
+                timestamp,
+                time_zone,
+            } => {
+                write!(f, "{} AT TIME ZONE '{}'", timestamp, time_zone)
+            }
         }
     }
 }

src/parser.rs

@@ -883,9 +883,17 @@ impl<'a> Parser<'a> {
     /// Parses an array expression `[ex1, ex2, ..]`
     /// if `named` is `true`, came from an expression like  `ARRAY[ex1, ex2]`
     pub fn parse_array_expr(&mut self, named: bool) -> Result<Expr, ParserError> {
-        let exprs = self.parse_comma_separated(Parser::parse_expr)?;
-        self.expect_token(&Token::RBracket)?;
-        Ok(Expr::Array(Array { elem: exprs, named }))
+        if self.peek_token() == Token::RBracket {
+            let _ = self.next_token();
+            Ok(Expr::Array(Array {
+                elem: vec![],
+                named,
+            }))
+        } else {
+            let exprs = self.parse_comma_separated(Parser::parse_expr)?;
+            self.expect_token(&Token::RBracket)?;
+            Ok(Expr::Array(Array { elem: exprs, named }))
+        }
     }
 
     /// Parse a SQL LISTAGG expression, e.g. `LISTAGG(...) WITHIN GROUP (ORDER BY ...)`.
@@ -1205,6 +1213,28 @@ impl<'a> Parser<'a> {
                         )
                     }
                 }
+                Keyword::AT => {
+                    // if self.parse_keyword(Keyword::TIME) {
+                    //     self.expect_keyword(Keyword::ZONE)?;
+                    if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) {
+                        let time_zone = self.next_token();
+                        match time_zone {
+                            Token::SingleQuotedString(time_zone) => {
+                                log::trace!("Peek token: {:?}", self.peek_token());
+                                Ok(Expr::AtTimeZone {
+                                    timestamp: Box::new(expr),
+                                    time_zone,
+                                })
+                            }
+                            tok => self.expected(
+                                "Expected Token::SingleQuotedString after AT TIME ZONE",
+                                tok,
+                            ),
+                        }
+                    } else {
+                        self.expected("Expected Token::Word after AT", tok)
+                    }
+                }
                 Keyword::NOT | Keyword::IN | Keyword::BETWEEN => {
                     self.prev_token();
                     let negated = self.parse_keyword(Keyword::NOT);
@@ -1350,15 +1380,32 @@ impl<'a> Parser<'a> {
     const UNARY_NOT_PREC: u8 = 15;
     const BETWEEN_PREC: u8 = 20;
     const PLUS_MINUS_PREC: u8 = 30;
+    const TIME_ZONE_PREC: u8 = 20;
 
     /// Get the precedence of the next token
     pub fn get_next_precedence(&self) -> Result<u8, ParserError> {
         let token = self.peek_token();
         debug!("get_next_precedence() {:?}", token);
+        let token_0 = self.peek_nth_token(0);
+        let token_1 = self.peek_nth_token(1);
+        let token_2 = self.peek_nth_token(2);
+        debug!("0: {token_0} 1: {token_1} 2: {token_2}");
         match token {
             Token::Word(w) if w.keyword == Keyword::OR => Ok(5),
             Token::Word(w) if w.keyword == Keyword::AND => Ok(10),
             Token::Word(w) if w.keyword == Keyword::XOR => Ok(24),
+
+            Token::Word(w) if w.keyword == Keyword::AT => {
+                match (self.peek_nth_token(1), self.peek_nth_token(2)) {
+                    (Token::Word(w), Token::Word(w2))
+                        if w.keyword == Keyword::TIME && w2.keyword == Keyword::ZONE =>
+                    {
+                        Ok(Self::TIME_ZONE_PREC)
+                    }
+                    _ => Ok(0),
+                }
+            }
+
             Token::Word(w) if w.keyword == Keyword::NOT => match self.peek_nth_token(1) {
                 // The precedence of NOT varies depending on keyword that
                 // follows it. If it is followed by IN, BETWEEN, or LIKE,

tests/sqlparser_common.rs

@@ -2871,6 +2871,65 @@ fn parse_literal_interval() {
     );
 }
 
+#[test]
+fn parse_at_timezone() {
+    let zero = Expr::Value(number("0"));
+    let sql = "SELECT FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00' FROM t";
+    let select = verified_only_select(sql);
+    assert_eq!(
+        &Expr::AtTimeZone {
+            timestamp: Box::new(Expr::Function(Function {
+                name: ObjectName(vec![Ident {
+                    value: "FROM_UNIXTIME".to_string(),
+                    quote_style: None
+                }]),
+                args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(zero.clone()))],
+                over: None,
+                distinct: false
+            })),
+            time_zone: "UTC-06:00".to_string()
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+
+    let sql = r#"SELECT DATE_FORMAT(FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00', '%Y-%m-%dT%H') AS "hour" FROM t"#;
+    let select = verified_only_select(sql);
+    assert_eq!(
+        &SelectItem::ExprWithAlias {
+            expr: Expr::Function(Function {
+                name: ObjectName(vec![Ident {
+                    value: "DATE_FORMAT".to_string(),
+                    quote_style: None,
+                },],),
+                args: vec![
+                    FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::AtTimeZone {
+                        timestamp: Box::new(Expr::Function(Function {
+                            name: ObjectName(vec![Ident {
+                                value: "FROM_UNIXTIME".to_string(),
+                                quote_style: None,
+                            },],),
+                            args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(zero,),),],
+                            over: None,
+                            distinct: false,
+                        },)),
+                        time_zone: "UTC-06:00".to_string(),
+                    },),),
+                    FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
+                        Value::SingleQuotedString("%Y-%m-%dT%H".to_string(),),
+                    ),),),
+                ],
+                over: None,
+                distinct: false,
+            },),
+            alias: Ident {
+                value: "hour".to_string(),
+                quote_style: Some('"',),
+            },
+        },
+        only(&select.projection),
+    );
+}
+
 #[test]
 fn parse_simple_math_expr_plus() {
     let sql = "SELECT a + b, 2 + a, 2.5 + a, a_f + b_f, 2 + a_f, 2.5 + a_f FROM c";

tests/sqlparser_postgres.rs

@@ -1228,6 +1228,16 @@ fn parse_array_index_expr() {
         },
         expr_from_projection(only(&select.projection)),
     );
+
+    let sql = "SELECT ARRAY[]";
+    let select = pg_and_generic().verified_only_select(sql);
+    assert_eq!(
+        &Expr::Array(sqlparser::ast::Array {
+            elem: vec![],
+            named: true
+        }),
+        expr_from_projection(only(&select.projection)),
+    );
 }
 
 #[test]