Merge pull request #228 from compareasiagroup/master

Fix for #227 (get_type() doesn't work for queries that use WITH)

1 files changed, 12 insertions, 0 deletions

diff --git a/sqlparse/sql.py b/sqlparse/sql.py
index c1111bb..9c0497a 100644
--- a/sqlparse/sql.py
+++ b/sqlparse/sql.py
@@ -487,6 +487,18 @@ class Statement(TokenList):
         elif first_token.ttype in (T.Keyword.DML, T.Keyword.DDL):
             return first_token.normalized
 
+        elif first_token.ttype == T.Keyword.CTE:
+            # The WITH keyword should be followed by either an Identifier or
+            # an IdentifierList containing the CTE definitions;  the actual
+            # DML keyword (e.g. SELECT, INSERT) will follow next.
+            idents = self.token_next(self.token_index(first_token), skip_ws=True)
+            if isinstance(idents, (Identifier, IdentifierList)):
+                dml_keyword = self.token_next(self.token_index(idents), skip_ws=True)
+                if dml_keyword.ttype == T.Keyword.DML:
+                    return dml_keyword.normalized
+            # Hmm, probably invalid syntax, so return unknown.
+            return 'UNKNOWN'
+
         return 'UNKNOWN'