Fix DoubleColon cast skipping AT TIME ZONE apache#1266 (apache#1267)

Author: demetribu

src/parser/mod.rs

@@ -1535,19 +1535,23 @@ impl<'a> Parser<'a> {
     pub fn parse_optional_cast_format(&mut self) -> Result<Option<CastFormat>, ParserError> {
         if self.parse_keyword(Keyword::FORMAT) {
             let value = self.parse_value()?;
-            if self.parse_keywords(&[Keyword::AT, Keyword::TIME, Keyword::ZONE]) {
-                Ok(Some(CastFormat::ValueAtTimeZone(
-                    value,
-                    self.parse_value()?,
-                )))
-            } else {
-                Ok(Some(CastFormat::Value(value)))
+            match self.parse_optional_time_zone()? {
+                Some(tz) => Ok(Some(CastFormat::ValueAtTimeZone(value, tz))),
+                None => Ok(Some(CastFormat::Value(value))),
             }
         } else {
             Ok(None)
         }
     }
 
+    pub fn parse_optional_time_zone(&mut self) -> Result<Option<Value>, ParserError> {
+        if self.parse_keywords(&[Keyword::AT, Keyword::TIME, Keyword::ZONE]) {
+            self.parse_value().map(Some)
+        } else {
+            Ok(None)
+        }
+    }
+
     /// mssql-like convert function
     fn parse_mssql_convert(&mut self) -> Result<Expr, ParserError> {
         self.expect_token(&Token::LParen)?;
@@ -2541,12 +2545,35 @@ impl<'a> Parser<'a> {
                 ),
             }
         } else if Token::DoubleColon == tok {
-            Ok(Expr::Cast {
+            let data_type = self.parse_data_type()?;
+
+            let cast_expr = Expr::Cast {
                 kind: CastKind::DoubleColon,
                 expr: Box::new(expr),
-                data_type: self.parse_data_type()?,
+                data_type: data_type.clone(),
                 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 {

tests/sqlparser_common.rs

@@ -7022,6 +7022,27 @@ fn parse_set_variable() {
     one_statement_parses_to("SET SOMETHING TO '1'", "SET SOMETHING = '1'");
 }
 
+#[test]
+fn parse_double_colon_cast_at_timezone() {
+    let sql = "SELECT '2001-01-01T00:00:00.000Z'::TIMESTAMP AT TIME ZONE 'Europe/Brussels' FROM t";
+    let select = verified_only_select(sql);
+
+    assert_eq!(
+        &Expr::AtTimeZone {
+            timestamp: Box::new(Expr::Cast {
+                kind: CastKind::DoubleColon,
+                expr: Box::new(Expr::Value(Value::SingleQuotedString(
+                    "2001-01-01T00:00:00.000Z".to_string()
+                ),)),
+                data_type: DataType::Timestamp(None, TimezoneInfo::None),
+                format: None
+            }),
+            time_zone: "Europe/Brussels".to_string()
+        },
+        expr_from_projection(only(&select.projection)),
+    );
+}
+
 #[test]
 fn parse_set_time_zone() {
     match verified_stmt("SET TIMEZONE = 'UTC'") {