From 70eb3d08c454ae1483d7c369ca106478b9892791 Mon Sep 17 00:00:00 2001 From: xiangjinwu <17769960+xiangjinwu@users.noreply.github.com> Date: Fri, 21 Jun 2024 18:18:15 +0800 Subject: [PATCH] feat(sqlparser): allow rhs of `AT TIME ZONE` to be non-literal (#17395) --- e2e_test/batch/functions/at_time_zone.slt.part | 16 ++++++++++++++++ src/frontend/src/binder/expr/mod.rs | 6 +++--- src/sqlparser/src/ast/mod.rs | 4 ++-- src/sqlparser/src/parser.rs | 17 +++++++++++------ src/sqlparser/tests/testdata/select.yaml | 13 ++++++++++++- 5 files changed, 44 insertions(+), 12 deletions(-) diff --git a/e2e_test/batch/functions/at_time_zone.slt.part b/e2e_test/batch/functions/at_time_zone.slt.part index e24bf7e41ef58..cfd89b3312965 100644 --- a/e2e_test/batch/functions/at_time_zone.slt.part +++ b/e2e_test/batch/functions/at_time_zone.slt.part @@ -21,3 +21,19 @@ query T select '2022-11-06 01:00:00'::timestamp AT TIME ZONE 'us/pacific'; ---- 2022-11-06 09:00:00+00:00 + +# non-literal zone +statement ok +create table t (local timestamp, tz varchar); + +statement ok +insert into t values ('2024-06-10 12:00:00', 'US/Pacific'), ('2024-06-10 13:00:00', 'Asia/Singapore'); + +query T +select local AT TIME ZONE tz from t order by 1; +---- +2024-06-10 05:00:00+00:00 +2024-06-10 19:00:00+00:00 + +statement ok +drop table t; diff --git a/src/frontend/src/binder/expr/mod.rs b/src/frontend/src/binder/expr/mod.rs index 0b8b50be0eab6..363e6f0738fef 100644 --- a/src/frontend/src/binder/expr/mod.rs +++ b/src/frontend/src/binder/expr/mod.rs @@ -175,7 +175,7 @@ impl Binder { Expr::AtTimeZone { timestamp, time_zone, - } => self.bind_at_time_zone(*timestamp, time_zone), + } => self.bind_at_time_zone(*timestamp, *time_zone), // special syntax for string Expr::Trim { expr, @@ -219,9 +219,9 @@ impl Binder { .into()) } - pub(super) fn bind_at_time_zone(&mut self, input: Expr, time_zone: String) -> Result { + pub(super) fn bind_at_time_zone(&mut self, input: Expr, time_zone: Expr) -> Result { let input = self.bind_expr_inner(input)?; - let time_zone = self.bind_string(time_zone)?.into(); + let time_zone = self.bind_expr_inner(time_zone)?; FunctionCall::new(ExprType::AtTimeZone, vec![input, time_zone]).map(Into::into) } diff --git a/src/sqlparser/src/ast/mod.rs b/src/sqlparser/src/ast/mod.rs index 49fddbfaa4b82..c27508e8e80d1 100644 --- a/src/sqlparser/src/ast/mod.rs +++ b/src/sqlparser/src/ast/mod.rs @@ -417,7 +417,7 @@ pub enum Expr { /// explicitly specified zone AtTimeZone { timestamp: Box, - time_zone: String, + time_zone: Box, }, /// `EXTRACT(DateTimeField FROM )` Extract { @@ -667,7 +667,7 @@ impl fmt::Display for Expr { Expr::AtTimeZone { timestamp, time_zone, - } => write!(f, "{} AT TIME ZONE '{}'", timestamp, time_zone), + } => write!(f, "{} AT TIME ZONE {}", timestamp, time_zone), Expr::Extract { field, expr } => write!(f, "EXTRACT({} FROM {})", field, expr), Expr::Collate { expr, collation } => write!(f, "{} COLLATE {}", expr, collation), Expr::Nested(ast) => write!(f, "({})", ast), diff --git a/src/sqlparser/src/parser.rs b/src/sqlparser/src/parser.rs index dbc79542949d2..7015bfac429f0 100644 --- a/src/sqlparser/src/parser.rs +++ b/src/sqlparser/src/parser.rs @@ -185,6 +185,8 @@ pub enum Precedence { PlusMinus, // 30 in upstream MulDiv, // 40 in upstream Exp, + At, + Collate, UnaryPosNeg, PostfixFactorial, Array, @@ -1396,11 +1398,14 @@ impl Parser<'_> { } } Keyword::AT => { - let time_zone = preceded( - (Keyword::TIME, Keyword::ZONE), - cut_err(Self::parse_literal_string), - ) - .parse_next(self)?; + assert_eq!(precedence, Precedence::At); + let time_zone = Box::new( + preceded( + (Keyword::TIME, Keyword::ZONE), + cut_err(|p: &mut Self| p.parse_subexpr(precedence)), + ) + .parse_next(self)?, + ); Ok(Expr::AtTimeZone { timestamp: Box::new(expr), time_zone, @@ -1657,7 +1662,7 @@ impl Parser<'_> { (Token::Word(w), Token::Word(w2)) if w.keyword == Keyword::TIME && w2.keyword == Keyword::ZONE => { - Ok(P::Other) + Ok(P::At) } _ => Ok(P::Zero), } diff --git a/src/sqlparser/tests/testdata/select.yaml b/src/sqlparser/tests/testdata/select.yaml index ad97f6568787b..781f65760258a 100644 --- a/src/sqlparser/tests/testdata/select.yaml +++ b/src/sqlparser/tests/testdata/select.yaml @@ -93,7 +93,18 @@ ^ - input: SELECT timestamp with time zone '2022-10-01 12:00:00Z' AT TIME ZONE 'US/Pacific' formatted_sql: SELECT TIMESTAMP WITH TIME ZONE '2022-10-01 12:00:00Z' AT TIME ZONE 'US/Pacific' - formatted_ast: 'Query(Query { with: None, body: Select(Select { distinct: All, projection: [UnnamedExpr(AtTimeZone { timestamp: TypedString { data_type: Timestamp(true), value: "2022-10-01 12:00:00Z" }, time_zone: "US/Pacific" })], from: [], lateral_views: [], selection: None, group_by: [], having: None }), order_by: [], limit: None, offset: None, fetch: None })' + formatted_ast: 'Query(Query { with: None, body: Select(Select { distinct: All, projection: [UnnamedExpr(AtTimeZone { timestamp: TypedString { data_type: Timestamp(true), value: "2022-10-01 12:00:00Z" }, time_zone: Value(SingleQuotedString("US/Pacific")) })], from: [], lateral_views: [], selection: None, group_by: [], having: None }), order_by: [], limit: None, offset: None, fetch: None })' +- input: SELECT timestamp with time zone '2022-10-01 12:00:00Z' AT TIME ZONE zone + formatted_sql: SELECT TIMESTAMP WITH TIME ZONE '2022-10-01 12:00:00Z' AT TIME ZONE zone + formatted_ast: 'Query(Query { with: None, body: Select(Select { distinct: All, projection: [UnnamedExpr(AtTimeZone { timestamp: TypedString { data_type: Timestamp(true), value: "2022-10-01 12:00:00Z" }, time_zone: Identifier(Ident { value: "zone", quote_style: None }) })], from: [], lateral_views: [], selection: None, group_by: [], having: None }), order_by: [], limit: None, offset: None, fetch: None })' +# https://www.postgresql.org/message-id/CADT4RqBPdbsZW7HS1jJP319TMRHs1hzUiP=iRJYR6UqgHCrgNQ@mail.gmail.com +- input: SELECT now() + INTERVAL '14 days' AT TIME ZONE 'UTC'; + formatted_sql: SELECT now() + INTERVAL '14 days' AT TIME ZONE 'UTC' + formatted_ast: 'Query(Query { with: None, body: Select(Select { distinct: All, projection: [UnnamedExpr(BinaryOp { left: Function(Function { name: ObjectName([Ident { value: "now", quote_style: None }]), args: [], variadic: false, over: None, distinct: false, order_by: [], filter: None, within_group: None }), op: Plus, right: AtTimeZone { timestamp: Value(Interval { value: "14 days", leading_field: None, leading_precision: None, last_field: None, fractional_seconds_precision: None }), time_zone: Value(SingleQuotedString("UTC")) } })], from: [], lateral_views: [], selection: None, group_by: [], having: None }), order_by: [], limit: None, offset: None, fetch: None })' +# https://github.com/sqlparser-rs/sqlparser-rs/issues/1266 +- input: SELECT c FROM t WHERE c >= '2019-03-27T22:00:00.000Z'::timestamp AT TIME ZONE 'Europe/Brussels'; + formatted_sql: SELECT c FROM t WHERE c >= CAST('2019-03-27T22:00:00.000Z' AS TIMESTAMP) AT TIME ZONE 'Europe/Brussels' + formatted_ast: 'Query(Query { with: None, body: Select(Select { distinct: All, projection: [UnnamedExpr(Identifier(Ident { value: "c", quote_style: None }))], from: [TableWithJoins { relation: Table { name: ObjectName([Ident { value: "t", quote_style: None }]), alias: None, as_of: None }, joins: [] }], lateral_views: [], selection: Some(BinaryOp { left: Identifier(Ident { value: "c", quote_style: None }), op: GtEq, right: AtTimeZone { timestamp: Cast { expr: Value(SingleQuotedString("2019-03-27T22:00:00.000Z")), data_type: Timestamp(false) }, time_zone: Value(SingleQuotedString("Europe/Brussels")) } }), group_by: [], having: None }), order_by: [], limit: None, offset: None, fetch: None })' - input: SELECT 0c6 error_msg: |- sql parser error: trailing junk after numeric literal at line 1, column 9