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

This commit is contained in:
Joey Hain 2024-05-23 10:30:05 -07:00 committed by GitHub
parent 9d15f7e9a9
commit d5faf3c54b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 65 additions and 53 deletions

View file

@ -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}")

View file

@ -2469,27 +2469,12 @@ 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());
self.expect_keywords(&[Keyword::TIME, Keyword::ZONE])?;
Ok(Expr::AtTimeZone {
timestamp: Box::new(expr),
time_zone,
time_zone: Box::new(self.parse_subexpr(precedence)?),
})
}
_ => self.expected(
"Expected Token::SingleQuotedString after AT TIME ZONE",
time_zone,
),
}
} else {
self.expected("Expected Token::Word after AT", tok)
}
}
Keyword::NOT
| Keyword::IN
| Keyword::BETWEEN
@ -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),
}

View file

@ -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)),
);

View file

@ -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
);
}