From f516b66a0e254af510b6b8a18510aad922d69701 Mon Sep 17 00:00:00 2001 From: Andrew Tipton Date: Wed, 2 Mar 2016 18:37:28 +0800 Subject: Ensure get_type() works for queries that use WITH. --- sqlparse/keywords.py | 2 +- sqlparse/sql.py | 12 ++++++++++++ sqlparse/tokens.py | 1 + 3 files changed, 14 insertions(+), 1 deletion(-) (limited to 'sqlparse') diff --git a/sqlparse/keywords.py b/sqlparse/keywords.py index b6c4246..dd08be0 100644 --- a/sqlparse/keywords.py +++ b/sqlparse/keywords.py @@ -486,7 +486,7 @@ KEYWORDS = { 'VOLATILE': tokens.Keyword, 'WHENEVER': tokens.Keyword, - 'WITH': tokens.Keyword, + 'WITH': tokens.Keyword.CTE, 'WITHOUT': tokens.Keyword, 'WORK': tokens.Keyword, 'WRITE': tokens.Keyword, 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' diff --git a/sqlparse/tokens.py b/sqlparse/tokens.py index 01a9b89..98fa8a6 100644 --- a/sqlparse/tokens.py +++ b/sqlparse/tokens.py @@ -75,6 +75,7 @@ Token.Number = Number # SQL specific tokens DML = Keyword.DML DDL = Keyword.DDL +CTE = Keyword.CTE Command = Keyword.Command Group = Token.Group -- cgit v1.2.1