feat(sqlparser): allow rhs of AT TIME ZONE to be non-literal (#17395)

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;

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<ExprImpl> {
+    pub(super) fn bind_at_time_zone(&mut self, input: Expr, time_zone: Expr) -> Result<ExprImpl> {
         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)
     }
 

src/sqlparser/src/ast/mod.rs

@@ -416,7 +416,7 @@ pub enum Expr {
     /// explicitly specified zone
     AtTimeZone {
         timestamp: Box<Expr>,
-        time_zone: String,
+        time_zone: Box<Expr>,
     },
     /// `EXTRACT(DateTimeField FROM <expr>)`
     Extract {
@@ -666,7 +666,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),

src/sqlparser/src/parser.rs

@@ -167,6 +167,8 @@ pub enum Precedence {
     PlusMinus, // 30 in upstream
     MulDiv,    // 40 in upstream
     Exp,
+    At,
+    Collate,
     UnaryPosNeg,
     PostfixFactorial,
     Array,
@@ -1539,17 +1541,12 @@ impl Parser {
                 }
                 Keyword::AT => {
                     if self.parse_keywords(&[Keyword::TIME, Keyword::ZONE]) {
-                        let token = self.next_token();
-                        match token.token {
-                            Token::SingleQuotedString(time_zone) => Ok(Expr::AtTimeZone {
-                                timestamp: Box::new(expr),
-                                time_zone,
-                            }),
-                            unexpected => self.expected(
-                                "Expected Token::SingleQuotedString after AT TIME ZONE",
-                                unexpected.with_location(token.location),
-                            ),
-                        }
+                        assert_eq!(precedence, Precedence::At);
+                        let time_zone = Box::new(self.parse_subexpr(precedence)?);
+                        Ok(Expr::AtTimeZone {
+                            timestamp: Box::new(expr),
+                            time_zone,
+                        })
                     } else {
                         self.expected("Expected Token::Word after AT", tok)
                     }
@@ -1807,7 +1804,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),
                 }

src/sqlparser/tests/testdata/select.yaml

@@ -76,7 +76,18 @@
   error_msg: 'sql parser error: syntax error at or near WHERE at line:1, column:16'
 - 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 8'
 - input: SELECT 1e6