Make tzcast grouping function less eager

2 files changed, 18 insertions, 2 deletions

diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
index 175ae8e..2fb0a4c 100644
--- a/sqlparse/engine/grouping.py
+++ b/sqlparse/engine/grouping.py
@@ -91,13 +91,20 @@ def group_tzcasts(tlist):
     def match(token):
         return token.ttype == T.Keyword.TZCast
 
-    def valid(token):
+    def valid_prev(token):
         return token is not None
 
+    def valid_next(token):
+        return token is not None and (
+            token.is_whitespace
+            or token.match(T.Keyword, 'AS')
+            or token.match(*sql.TypedLiteral.M_CLOSE)
+        )
+
     def post(tlist, pidx, tidx, nidx):
         return pidx, nidx
 
-    _group(tlist, sql.Identifier, match, valid, valid, post)
+    _group(tlist, sql.Identifier, match, valid_prev, valid_next, post)
 
 
 def group_typed_literal(tlist):
diff --git a/tests/test_regressions.py b/tests/test_regressions.py
index 38d1840..4ffc69f 100644
--- a/tests/test_regressions.py
+++ b/tests/test_regressions.py
@@ -401,6 +401,15 @@ def test_issue489_tzcasts():
     assert p.tokens[-1].get_alias() == 'foo'
 
 
+def test_issue562_tzcasts():
+    # Test that whitespace between 'from' and 'bar' is retained
+    formatted = sqlparse.format(
+        'SELECT f(HOUR from bar AT TIME ZONE \'UTC\') from foo', reindent=True
+    )
+    assert formatted == \
+           'SELECT f(HOUR\n         from bar AT TIME ZONE \'UTC\')\nfrom foo'
+
+
 def test_as_in_parentheses_indents():
     # did raise NoneType has no attribute is_group in _process_parentheses
     formatted = sqlparse.format('(as foo)', reindent=True)