Add support for time zone casts (fixes andialbrecht#489).

andialbrecht · andialbrecht · commit 7168e5ce6711 · 2019-05-16T07:56:57.000+02:00
diff --git a/CHANGELOG b/CHANGELOG
@@ -4,6 +4,7 @@ Development Version
 Enhancements
 
 * Add HQL keywords (pr475, by matwalk).
+* Add support for time zone casts (issue489).
 
 Bug Fixes
 
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
@@ -88,6 +88,19 @@ def post(tlist, pidx, tidx, nidx):
     _group(tlist, sql.Identifier, match, valid_prev, valid_next, post)
 
 
+def group_tzcasts(tlist):
+    def match(token):
+        return token.ttype == T.Keyword.TZCast
+
+    def valid(token):
+        return token is not None
+
+    def post(tlist, pidx, tidx, nidx):
+        return pidx, nidx
+
+    _group(tlist, sql.Identifier, match, valid, valid, post)
+
+
 def group_period(tlist):
     def match(token):
         return token.match(T.Punctuation, '.')
@@ -358,6 +371,7 @@ def group(stmt):
         group_identifier,
         group_order,
         group_typecasts,
+        group_tzcasts,
         group_operator,
         group_comparison,
         group_as,
diff --git a/sqlparse/keywords.py b/sqlparse/keywords.py
@@ -85,7 +85,7 @@ def is_keyword(value):
         (r'GROUP\s+BY\b', tokens.Keyword),
         (r'ORDER\s+BY\b', tokens.Keyword),
         (r'(LATERAL\s+VIEW\s+)(EXPLODE|INLINE|PARSE_URL_TUPLE|POSEXPLODE|STACK)\b', tokens.Keyword),
-
+        (r"(AT|WITH')\s+TIME\s+ZONE\s+'[^']+'", tokens.Keyword.TZCast),
         (r'[0-9_A-Z&Agrave;-&Uuml;][_$#\w]*', is_keyword),
 
         (r'[;:()\[\],\.]', tokens.Punctuation),
diff --git a/tests/test_regressions.py b/tests/test_regressions.py
@@ -400,3 +400,9 @@ def test_issue485_split_multi():
 select 2;
 );'''
     assert len(sqlparse.split(p_sql)) == 1
+
+
+def test_issue489_tzcasts():
+    p = sqlparse.parse('select bar at time zone \'UTC\' as foo')[0]
+    assert p.tokens[-1].has_alias() is True
+    assert p.tokens[-1].get_alias() == 'foo'
diff --git a/tests/test_tokenize.py b/tests/test_tokenize.py
@@ -204,6 +204,15 @@ def test_parse_order_by():
     assert p.tokens[0].ttype is T.Keyword
 
 
+@pytest.mark.parametrize('s', (
+    'AT TIME ZONE \'UTC\'',
+))
+def test_parse_tzcast(s):
+    p = sqlparse.parse(s)[0]
+    assert len(p.tokens) == 1
+    assert p.tokens[0].ttype == T.Keyword.TZCast
+
+
 def test_cli_commands():
     p = sqlparse.parse('\\copy')[0]
     assert len(p.tokens) == 1
