Skip to content

Commit

Permalink
Add support for AT TIME ZONE (#539)
Browse files Browse the repository at this point in the history
* Support for empty array literals

* Added support for AT TIME ZONE

Co-authored-by: Chris Allen <chrisa@indeed.com>
  • Loading branch information
bitemyapp and Chris Allen committed Jul 18, 2022
1 parent 4706d8b commit 7cbbd91
Show file tree
Hide file tree
Showing 4 changed files with 130 additions and 3 deletions.
11 changes: 11 additions & 0 deletions src/ast/mod.rs
Expand Up @@ -285,6 +285,11 @@ pub enum Expr {
expr: Box<Expr>,
data_type: DataType,
},
/// 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,
},
/// EXTRACT(DateTimeField FROM <expr>)
Extract {
field: DateTimeField,
Expand Down Expand Up @@ -562,6 +567,12 @@ impl fmt::Display for Expr {
Expr::CompositeAccess { expr, key } => {
write!(f, "{}.{}", expr, key)
}
Expr::AtTimeZone {
timestamp,
time_zone,
} => {
write!(f, "{} AT TIME ZONE '{}'", timestamp, time_zone)
}
}
}
}
Expand Down
53 changes: 50 additions & 3 deletions src/parser.rs
Expand Up @@ -883,9 +883,17 @@ impl<'a> Parser<'a> {
/// Parses an array expression `[ex1, ex2, ..]`
/// if `named` is `true`, came from an expression like `ARRAY[ex1, ex2]`
pub fn parse_array_expr(&mut self, named: bool) -> Result<Expr, ParserError> {
let exprs = self.parse_comma_separated(Parser::parse_expr)?;
self.expect_token(&Token::RBracket)?;
Ok(Expr::Array(Array { elem: exprs, named }))
if self.peek_token() == Token::RBracket {
let _ = self.next_token();
Ok(Expr::Array(Array {
elem: vec![],
named,
}))
} else {
let exprs = self.parse_comma_separated(Parser::parse_expr)?;
self.expect_token(&Token::RBracket)?;
Ok(Expr::Array(Array { elem: exprs, named }))
}
}

/// Parse a SQL LISTAGG expression, e.g. `LISTAGG(...) WITHIN GROUP (ORDER BY ...)`.
Expand Down Expand Up @@ -1205,6 +1213,28 @@ 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::SingleQuotedString(time_zone) => {
log::trace!("Peek token: {:?}", self.peek_token());
Ok(Expr::AtTimeZone {
timestamp: Box::new(expr),
time_zone,
})
}
tok => self.expected(
"Expected Token::SingleQuotedString after AT TIME ZONE",
tok,
),
}
} else {
self.expected("Expected Token::Word after AT", tok)
}
}
Keyword::NOT | Keyword::IN | Keyword::BETWEEN => {
self.prev_token();
let negated = self.parse_keyword(Keyword::NOT);
Expand Down Expand Up @@ -1350,15 +1380,32 @@ impl<'a> Parser<'a> {
const UNARY_NOT_PREC: u8 = 15;
const BETWEEN_PREC: u8 = 20;
const PLUS_MINUS_PREC: u8 = 30;
const TIME_ZONE_PREC: u8 = 20;

/// Get the precedence of the next token
pub fn get_next_precedence(&self) -> Result<u8, ParserError> {
let token = self.peek_token();
debug!("get_next_precedence() {:?}", token);
let token_0 = self.peek_nth_token(0);
let token_1 = self.peek_nth_token(1);
let token_2 = self.peek_nth_token(2);
debug!("0: {token_0} 1: {token_1} 2: {token_2}");
match token {
Token::Word(w) if w.keyword == Keyword::OR => Ok(5),
Token::Word(w) if w.keyword == Keyword::AND => Ok(10),
Token::Word(w) if w.keyword == Keyword::XOR => Ok(24),

Token::Word(w) if w.keyword == Keyword::AT => {
match (self.peek_nth_token(1), self.peek_nth_token(2)) {
(Token::Word(w), Token::Word(w2))
if w.keyword == Keyword::TIME && w2.keyword == Keyword::ZONE =>
{
Ok(Self::TIME_ZONE_PREC)
}
_ => Ok(0),
}
}

Token::Word(w) if w.keyword == Keyword::NOT => match self.peek_nth_token(1) {
// The precedence of NOT varies depending on keyword that
// follows it. If it is followed by IN, BETWEEN, or LIKE,
Expand Down
59 changes: 59 additions & 0 deletions tests/sqlparser_common.rs
Expand Up @@ -2871,6 +2871,65 @@ fn parse_literal_interval() {
);
}

#[test]
fn parse_at_timezone() {
let zero = Expr::Value(number("0"));
let sql = "SELECT FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00' FROM t";
let select = verified_only_select(sql);
assert_eq!(
&Expr::AtTimeZone {
timestamp: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident {
value: "FROM_UNIXTIME".to_string(),
quote_style: None
}]),
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(zero.clone()))],
over: None,
distinct: false
})),
time_zone: "UTC-06:00".to_string()
},
expr_from_projection(only(&select.projection)),
);

let sql = r#"SELECT DATE_FORMAT(FROM_UNIXTIME(0) AT TIME ZONE 'UTC-06:00', '%Y-%m-%dT%H') AS "hour" FROM t"#;
let select = verified_only_select(sql);
assert_eq!(
&SelectItem::ExprWithAlias {
expr: Expr::Function(Function {
name: ObjectName(vec![Ident {
value: "DATE_FORMAT".to_string(),
quote_style: None,
},],),
args: vec![
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::AtTimeZone {
timestamp: Box::new(Expr::Function(Function {
name: ObjectName(vec![Ident {
value: "FROM_UNIXTIME".to_string(),
quote_style: None,
},],),
args: vec![FunctionArg::Unnamed(FunctionArgExpr::Expr(zero,),),],
over: None,
distinct: false,
},)),
time_zone: "UTC-06:00".to_string(),
},),),
FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(
Value::SingleQuotedString("%Y-%m-%dT%H".to_string(),),
),),),
],
over: None,
distinct: false,
},),
alias: Ident {
value: "hour".to_string(),
quote_style: Some('"',),
},
},
only(&select.projection),
);
}

#[test]
fn parse_simple_math_expr_plus() {
let sql = "SELECT a + b, 2 + a, 2.5 + a, a_f + b_f, 2 + a_f, 2.5 + a_f FROM c";
Expand Down
10 changes: 10 additions & 0 deletions tests/sqlparser_postgres.rs
Expand Up @@ -1228,6 +1228,16 @@ fn parse_array_index_expr() {
},
expr_from_projection(only(&select.projection)),
);

let sql = "SELECT ARRAY[]";
let select = pg_and_generic().verified_only_select(sql);
assert_eq!(
&Expr::Array(sqlparser::ast::Array {
elem: vec![],
named: true
}),
expr_from_projection(only(&select.projection)),
);
}

#[test]
Expand Down

0 comments on commit 7cbbd91

Please sign in to comment.