[grouping] group_as() no longer groups AS CTE

liulk · andialbrecht · commit a301b79c042c · 2020-02-02T21:21:09.000+01:00
This patch changes the grouping of AS so that:

  Foo AS WITH bar AS 1 SELECT 2

with no longer be grouped as:

  [Identifier[Foo, AS, WITH, Identifier[Bar AS 1]], SELECT, 2]

but will be grouped as:

  [Identifier[Foo], AS, WITH, Identifier[Bar AS 1], SELECT, 2]

This fixes the parsing of CREATE TABLE new_table AS WITH ... so the
rest of the tokens after AS are parsed the same as a bare WITH.
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
@@ -164,7 +164,7 @@ def valid_prev(token):
         return token.normalized == 'NULL' or not token.is_keyword
 
     def valid_next(token):
-        ttypes = T.DML, T.DDL
+        ttypes = T.DML, T.DDL, T.CTE
         return not imt(token, t=ttypes) and token is not None
 
     def post(tlist, pidx, tidx, nidx):
diff --git a/tests/test_grouping.py b/tests/test_grouping.py
@@ -632,3 +632,11 @@ def test_aliased_literal_without_as():
     p = sqlparse.parse('1 foo')[0].tokens
     assert len(p) == 1
     assert p[0].get_alias() == 'foo'
+
+
+def test_grouping_as_cte():
+    p = sqlparse.parse('foo AS WITH apple AS 1, banana AS 2')[0].tokens
+    assert len(p) > 4
+    assert p[0].get_alias() is None
+    assert p[2].value == 'AS'
+    assert p[4].value == 'WITH'
