Add support for AT TIME ZONE expressions (#1196)

* Add support for AT TIME ZONE expressions

* adding tests

* Fixing imports

Author: tomershay

src/main/java/net/sf/jsqlparser/expression/ExpressionVisitor.java

@@ -163,4 +163,6 @@ public interface ExpressionVisitor {
     void visit(VariableAssignment aThis);
 
     void visit(XMLSerializeExpr aThis);
+
+    void visit(TimezoneExpression aThis);
 }

src/main/java/net/sf/jsqlparser/expression/ExpressionVisitorAdapter.java

@@ -576,4 +576,9 @@ public void visit(XMLSerializeExpr expr) {
             elm.getExpression().accept(this);
         }
     }
+
+    @Override
+    public void visit(TimezoneExpression expr) {
+        expr.getLeftExpression().accept(this);
+    }
 }

src/main/java/net/sf/jsqlparser/expression/TimezoneExpression.java

@@ -0,0 +1,52 @@
+/*-
+ * #%L
+ * JSQLParser library
+ * %%
+ * Copyright (C) 2004 - 2019 JSQLParser
+ * %%
+ * Dual licensed under GNU LGPL 2.1 or Apache License 2.0
+ * #L%
+ */
+package net.sf.jsqlparser.expression;
+
+import net.sf.jsqlparser.parser.ASTNodeAccessImpl;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class TimezoneExpression extends ASTNodeAccessImpl implements Expression {
+
+    private Expression leftExpression;
+    private ArrayList<String> timezoneExpressions = new ArrayList<>();
+
+    public Expression getLeftExpression() {
+        return leftExpression;
+    }
+
+    public void setLeftExpression(Expression expression) {
+        leftExpression = expression;
+    }
+
+    @Override
+    public void accept(ExpressionVisitor expressionVisitor) {
+        expressionVisitor.visit(this);
+    }
+
+    public List<String> getTimezoneExpressions() {
+        return timezoneExpressions;
+    }
+
+    public void addTimezoneExpression(String timezoneExpr) {
+        this.timezoneExpressions.add(timezoneExpr);
+    }
+
+    @Override
+    public String toString() {
+        String returnValue = getLeftExpression().toString();
+        for (String expr : timezoneExpressions) {
+            returnValue += " AT TIME ZONE " + expr;
+        }
+
+        return returnValue;
+    }
+}

src/main/java/net/sf/jsqlparser/util/TablesNamesFinder.java

@@ -11,6 +11,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+
 import net.sf.jsqlparser.expression.AllComparisonExpression;
 import net.sf.jsqlparser.expression.AnalyticExpression;
 import net.sf.jsqlparser.expression.AnyComparisonExpression;
@@ -48,6 +49,7 @@
 import net.sf.jsqlparser.expression.TimeKeyExpression;
 import net.sf.jsqlparser.expression.TimeValue;
 import net.sf.jsqlparser.expression.TimestampValue;
+import net.sf.jsqlparser.expression.TimezoneExpression;
 import net.sf.jsqlparser.expression.UserVariable;
 import net.sf.jsqlparser.expression.ValueListExpression;
 import net.sf.jsqlparser.expression.VariableAssignment;
@@ -928,4 +930,9 @@ public void visit(CreateSynonym createSynonym) {
     private static <T> void throwUnsupported(T type){
         throw new UnsupportedOperationException(String.format("Finding tables from %s is not supported", type.getClass().getSimpleName()));
     }
+
+    @Override
+    public void visit(TimezoneExpression aThis) {
+        aThis.getLeftExpression().accept(this);
+    }
 }

src/main/java/net/sf/jsqlparser/util/deparser/ExpressionDeParser.java

@@ -51,6 +51,7 @@
 import net.sf.jsqlparser.expression.TimeKeyExpression;
 import net.sf.jsqlparser.expression.TimeValue;
 import net.sf.jsqlparser.expression.TimestampValue;
+import net.sf.jsqlparser.expression.TimezoneExpression;
 import net.sf.jsqlparser.expression.UserVariable;
 import net.sf.jsqlparser.expression.ValueListExpression;
 import net.sf.jsqlparser.expression.VariableAssignment;
@@ -989,4 +990,13 @@ public void visit(XMLSerializeExpr expr) {
         }
         buffer.append(") AS ").append(expr.getDataType()).append(")");
     }
+
+    @Override
+    public void visit(TimezoneExpression var) {
+        var.getLeftExpression().accept(this);
+
+        for (String expr : var.getTimezoneExpressions()) {
+            buffer.append(" AT TIME ZONE " + expr);
+        }
+    }
 }

src/main/java/net/sf/jsqlparser/util/validation/validator/ExpressionValidator.java

@@ -46,6 +46,7 @@
 import net.sf.jsqlparser.expression.TimeKeyExpression;
 import net.sf.jsqlparser.expression.TimeValue;
 import net.sf.jsqlparser.expression.TimestampValue;
+import net.sf.jsqlparser.expression.TimezoneExpression;
 import net.sf.jsqlparser.expression.UserVariable;
 import net.sf.jsqlparser.expression.ValueListExpression;
 import net.sf.jsqlparser.expression.VariableAssignment;
@@ -563,6 +564,11 @@ public void visit(VariableAssignment a) {
         }
     }
 
+    @Override
+    public void visit(TimezoneExpression a) {
+        validateOptionalExpression(a.getLeftExpression());
+    }
+
     @Override
     public void visit(XMLSerializeExpr xml) {
         // TODO this feature seams very close to a jsqlparser-user usecase

src/main/jjtree/net/sf/jsqlparser/parser/JSqlParserCC.jjt

@@ -133,6 +133,7 @@ TOKEN: /* SQL Keywords. prefixed with K_ to avoid name clashes */
 |   <K_APPLY:"APPLY">
 |   <K_ARRAY_LITERAL: "ARRAY" >
 |   <K_AS: "AS">
+|   <K_AT: "AT">
 |   <K_ASC:"ASC">
 |   <K_AUTHORIZATION:"AUTHORIZATION">
 |   <K_BEGIN:"BEGIN">
@@ -3301,6 +3302,7 @@ Expression PrimaryExpression() #PrimaryExpression:
 {
     Expression retval = null;
     CastExpression castExpr = null;
+    TimezoneExpression timezoneExpr = null;
     Token token = null;
     Token sign = null;
     String tmp = "";
@@ -3399,6 +3401,20 @@ Expression PrimaryExpression() #PrimaryExpression:
         retval=castExpr;
     } )*
 
+    ( <K_AT> <K_DATETIMELITERAL> <K_ZONE> token=<S_CHAR_LITERAL> {
+        if (timezoneExpr == null)
+            timezoneExpr = new TimezoneExpression();
+
+        timezoneExpr.addTimezoneExpression(token.image);
+    } )*
+
+    {
+        if (timezoneExpr != null && !timezoneExpr.getTimezoneExpressions().isEmpty()) {
+            timezoneExpr.setLeftExpression(retval);
+            retval=timezoneExpr;
+        }
+    }
+
     {
         if (sign != null) {
             retval = new SignedExpression(sign.image.charAt(0), retval);

src/test/java/net/sf/jsqlparser/expression/ExpressionVisitorAdapterTest.java

@@ -208,4 +208,10 @@ public void testAnalyticFunctionWithoutExpression502() throws JSQLParserExceptio
         expr.accept(adapter);
     }
 
+    @Test
+    public void testAtTimeZoneExpression() throws JSQLParserException {
+        Expression expr = CCJSqlParserUtil.parseExpression("DATE(date1 AT TIME ZONE 'UTC' AT TIME ZONE 'australia/sydney')");
+        ExpressionVisitorAdapter adapter = new ExpressionVisitorAdapter();
+        expr.accept(adapter);
+    }
 }

src/test/java/net/sf/jsqlparser/statement/select/SelectTest.java

@@ -882,6 +882,18 @@ public void testSelectItems() throws JSQLParserException {
         assertStatementCanBeDeparsedAs(select, statement);
     }
 
+    @Test
+    public void testTimezoneExpression() throws JSQLParserException {
+        String stmt = "SELECT creation_date AT TIME ZONE 'UTC'";
+        assertSqlCanBeParsedAndDeparsed(stmt);
+    }
+
+    @Test
+    public void testTimezoneExpressionWithTwoTransformations() throws JSQLParserException {
+        String stmt = "SELECT DATE(date1 AT TIME ZONE 'UTC' AT TIME ZONE 'australia/sydney') AS another_date";
+        assertSqlCanBeParsedAndDeparsed(stmt);
+    }
+
     @Test
     public void testUnionWithOrderByAndLimitAndNoBrackets() throws JSQLParserException {
         String stmt = "SELECT id FROM table1 UNION SELECT id FROM table2 ORDER BY id ASC LIMIT 55";

src/test/java/net/sf/jsqlparser/util/TablesNamesFinderTest.java

@@ -642,4 +642,14 @@ public void testNPEIssue1009() throws JSQLParserException {
 
         assertThat(tablesNamesFinder.getTableList(stmt)).containsExactly("biz_fund_info");
     }
+
+    @Test
+    public void testAtTimeZoneExpression() throws JSQLParserException {
+        String sql = "SELECT DATE(date1 AT TIME ZONE 'UTC' AT TIME ZONE 'australia/sydney') AS another_date FROM mytbl";
+        Statement stmt = CCJSqlParserUtil.parse(sql);
+        TablesNamesFinder tablesNamesFinder = new TablesNamesFinder();
+        List<String> tableList = tablesNamesFinder.getTableList(stmt);
+        assertEquals(1, tableList.size());
+        assertTrue(tableList.contains("mytbl"));
+    }
 }