Support expression in AT TIME ZONE and fix precedence (apache#1272)

Author: jmhain

src/ast/mod.rs

@@ -584,7 +584,7 @@ pub enum Expr {
     /// 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,
+        time_zone: Box<Expr>,
     },
     /// Extract a field from a timestamp e.g. `EXTRACT(MONTH FROM foo)`
     ///
@@ -1270,7 +1270,7 @@ impl fmt::Display for Expr {
                 timestamp,
                 time_zone,
             } => {
-                write!(f, "{timestamp} AT TIME ZONE '{time_zone}'")
+                write!(f, "{timestamp} AT TIME ZONE {time_zone}")
             }
             Expr::Interval(interval) => {
                 write!(f, "{interval}")

src/parser/mod.rs

@@ -2469,26 +2469,11 @@ 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 {
-                            Token::SingleQuotedString(time_zone) => {
-                                log::trace!("Peek token: {:?}", self.peek_token());
-                                Ok(Expr::AtTimeZone {
-                                    timestamp: Box::new(expr),
-                                    time_zone,
-                                })
-                            }
-                            _ => self.expected(
-                                "Expected Token::SingleQuotedString after AT TIME ZONE",
-                                time_zone,
-                            ),
-                        }
-                    } else {
-                        self.expected("Expected Token::Word after AT", tok)
-                    }
+                    self.expect_keywords(&[Keyword::TIME, Keyword::ZONE])?;
+                    Ok(Expr::AtTimeZone {
+                        timestamp: Box::new(expr),
+                        time_zone: Box::new(self.parse_subexpr(precedence)?),
+                    })
                 }
                 Keyword::NOT
                 | Keyword::IN
@@ -2545,35 +2530,12 @@ impl<'a> Parser<'a> {
                 ),
             }
         } else if Token::DoubleColon == tok {
-            let data_type = self.parse_data_type()?;
-
-            let cast_expr = Expr::Cast {
+            Ok(Expr::Cast {
                 kind: CastKind::DoubleColon,
                 expr: Box::new(expr),
-                data_type: data_type.clone(),
+                data_type: self.parse_data_type()?,
                 format: None,
-            };
-
-            match data_type {
-                DataType::Date
-                | DataType::Datetime(_)
-                | DataType::Timestamp(_, _)
-                | DataType::Time(_, _) => {
-                    let value = self.parse_optional_time_zone()?;
-                    match value {
-                        Some(Value::SingleQuotedString(tz)) => Ok(Expr::AtTimeZone {
-                            timestamp: Box::new(cast_expr),
-                            time_zone: tz,
-                        }),
-                        None => Ok(cast_expr),
-                        _ => Err(ParserError::ParserError(format!(
-                            "Expected Token::SingleQuotedString after AT TIME ZONE, but found: {}",
-                            value.unwrap()
-                        ))),
-                    }
-                }
-                _ => Ok(cast_expr),
-            }
+            })
         } else if Token::ExclamationMark == tok {
             // PostgreSQL factorial operation
             Ok(Expr::UnaryOp {
@@ -2784,10 +2746,14 @@ impl<'a> Parser<'a> {
 
     // use https://www.postgresql.org/docs/7.0/operators.htm#AEN2026 as a reference
     // higher number = higher precedence
+    //
+    // NOTE: The pg documentation is incomplete, e.g. the AT TIME ZONE operator
+    //       actually has higher precedence than addition.
+    //       See https://postgrespro.com/list/thread-id/2673331.
+    const AT_TZ_PREC: u8 = 41;
     const MUL_DIV_MOD_OP_PREC: u8 = 40;
     const PLUS_MINUS_PREC: u8 = 30;
     const XOR_PREC: u8 = 24;
-    const TIME_ZONE_PREC: u8 = 20;
     const BETWEEN_PREC: u8 = 20;
     const LIKE_PREC: u8 = 19;
     const IS_PREC: u8 = 17;
@@ -2817,7 +2783,7 @@ impl<'a> Parser<'a> {
                     (Token::Word(w), Token::Word(w2))
                         if w.keyword == Keyword::TIME && w2.keyword == Keyword::ZONE =>
                     {
-                        Ok(Self::TIME_ZONE_PREC)
+                        Ok(Self::AT_TZ_PREC)
                     }
                     _ => Ok(0),
                 }

tests/sqlparser_common.rs

@@ -4995,7 +4995,9 @@ fn parse_at_timezone() {
     assert_eq!(
         &Expr::AtTimeZone {
             timestamp: Box::new(call("FROM_UNIXTIME", [zero.clone()])),
-            time_zone: "UTC-06:00".to_string(),
+            time_zone: Box::new(Expr::Value(Value::SingleQuotedString(
+                "UTC-06:00".to_string()
+            ))),
         },
         expr_from_projection(only(&select.projection)),
     );
@@ -5009,7 +5011,9 @@ fn parse_at_timezone() {
                 [
                     Expr::AtTimeZone {
                         timestamp: Box::new(call("FROM_UNIXTIME", [zero])),
-                        time_zone: "UTC-06:00".to_string(),
+                        time_zone: Box::new(Expr::Value(Value::SingleQuotedString(
+                            "UTC-06:00".to_string()
+                        ))),
                     },
                     Expr::Value(Value::SingleQuotedString("%Y-%m-%dT%H".to_string()),)
                 ]
@@ -7037,7 +7041,9 @@ fn parse_double_colon_cast_at_timezone() {
                 data_type: DataType::Timestamp(None, TimezoneInfo::None),
                 format: None
             }),
-            time_zone: "Europe/Brussels".to_string()
+            time_zone: Box::new(Expr::Value(Value::SingleQuotedString(
+                "Europe/Brussels".to_string()
+            ))),
         },
         expr_from_projection(only(&select.projection)),
     );

tests/sqlparser_postgres.rs

@@ -3882,3 +3882,43 @@ fn parse_mat_cte() {
     let sql2 = r#"WITH cte AS NOT MATERIALIZED (SELECT id FROM accounts) SELECT id FROM cte"#;
     pg().verified_stmt(sql2);
 }
+
+#[test]
+fn parse_at_time_zone() {
+    pg_and_generic().verified_expr("CURRENT_TIMESTAMP AT TIME ZONE tz");
+    pg_and_generic().verified_expr("CURRENT_TIMESTAMP AT TIME ZONE ('America/' || 'Los_Angeles')");
+
+    // check precedence
+    let expr = Expr::BinaryOp {
+        left: Box::new(Expr::AtTimeZone {
+            timestamp: Box::new(Expr::TypedString {
+                data_type: DataType::Timestamp(None, TimezoneInfo::None),
+                value: "2001-09-28 01:00".to_owned(),
+            }),
+            time_zone: Box::new(Expr::Cast {
+                kind: CastKind::DoubleColon,
+                expr: Box::new(Expr::Value(Value::SingleQuotedString(
+                    "America/Los_Angeles".to_owned(),
+                ))),
+                data_type: DataType::Text,
+                format: None,
+            }),
+        }),
+        op: BinaryOperator::Plus,
+        right: Box::new(Expr::Interval(Interval {
+            value: Box::new(Expr::Value(Value::SingleQuotedString(
+                "23 hours".to_owned(),
+            ))),
+            leading_field: None,
+            leading_precision: None,
+            last_field: None,
+            fractional_seconds_precision: None,
+        })),
+    };
+    pretty_assertions::assert_eq!(
+        pg_and_generic().verified_expr(
+            "TIMESTAMP '2001-09-28 01:00' AT TIME ZONE 'America/Los_Angeles'::TEXT + INTERVAL '23 hours'",
+        ),
+        expr
+    );
+}