Added support for AT TIME ZONE

sqlparser-rs · Jul 12, 2022 · f2a33a7 · f2a33a7
1 parent 6f09301
commit f2a33a7
Show file tree

Hide file tree

Showing 3 changed files with 123 additions and 0 deletions.
diff --git a/src/ast/mod.rs b/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,9 @@ 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)
+            }
         }
     }
 }

diff --git a/src/parser.rs b/src/parser.rs
@@ -1213,6 +1213,22 @@ 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);
@@ -1358,15 +1374,27 @@ 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,

diff --git a/tests/sqlparser_common.rs b/tests/sqlparser_common.rs
@@ -2790,6 +2790,93 @@ 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";