Group ROLE as an identifier when used as a column name (issue798)

ROLE is a non-reserved keyword, so it can also be a column name, but the
lexer always tokenized it as a keyword and `group_identifier` only wraps
Name/String.Symbol tokens. As a result `SELECT a, role, b` left `role` as a
bare keyword while `a` and `b` became Identifiers, so `get_identifiers()`
returned it inconsistently.

Simply dropping ROLE from the keyword list is not viable: it would make
`CREATE ROLE r` parse as an identifier aliased to the role name. Instead,
add `group_identifier_role`, which wraps ROLE in an Identifier only when it
sits in an identifier list (flanked by a comma). Outside a list -- e.g.
`CREATE ROLE`, `SET ROLE`, `DROP ROLE` -- there is no adjacent comma, so it
keeps its keyword role. This mirrors the existing `('null', 'role')`
special-case already present in `group_identifier_list`.

Author: adarshsm

CHANGELOG

@@ -61,6 +61,10 @@ Bug Fixes
   when it is followed by a float literal written without a leading zero, so
   that `x BETWEEN .03 AND .06` parses its bounds as numbers (issue601, pr868
   by deepakganesh78).
+* Group `ROLE` as an identifier when it is used as a column name in an
+  identifier list (e.g. `SELECT a, role, b`), so it appears alongside the
+  other columns in `get_identifiers()`, while keeping it a keyword elsewhere
+  such as `CREATE ROLE` (issue798, pr872 by adarshsm).
 * Fix a late-binding closure bug in `TokenList.token_not_matching`.
 
 Other

sqlparse/engine/grouping.py

@@ -254,6 +254,27 @@ def group_identifier(tlist):
         tidx, token = tlist.token_next_by(t=ttypes, idx=tidx)
 
 
+@recurse(sql.Identifier)
+def group_identifier_role(tlist):
+    # ROLE is a non-reserved keyword, so it may also be a column name. When it
+    # appears as a member of an identifier list (i.e. flanked by a comma) wrap
+    # it in an Identifier, matching the neighbouring plain names, so consumers
+    # such as ``get_identifiers()`` see a uniform Identifier rather than a bare
+    # keyword token. Outside of a list (e.g. ``CREATE ROLE r``) there is no
+    # adjacent comma, so it keeps its keyword role. See issue #798.
+    m_role = T.Keyword, 'ROLE'
+    tidx, token = tlist.token_next_by(m=m_role)
+    while token:
+        _, prev_ = tlist.token_prev(tidx)
+        _, next_ = tlist.token_next(tidx)
+        in_list = (prev_ is not None and prev_.match(T.Punctuation, ',')) or (
+            next_ is not None and next_.match(T.Punctuation, ',')
+        )
+        if in_list:
+            tlist.group_tokens(sql.Identifier, tidx, tidx)
+        tidx, token = tlist.token_next_by(m=m_role, idx=tidx)
+
+
 @recurse(sql.Over)
 def group_over(tlist):
     tidx, token = tlist.token_next_by(m=sql.Over.M_OPEN)
@@ -457,6 +478,7 @@ def group(stmt):
         group_period,
         group_arrays,
         group_identifier,
+        group_identifier_role,
         group_order,
         group_typecasts,
         group_tzcasts,

tests/test_grouping.py

@@ -241,6 +241,26 @@ def test_grouping_identifier_list_with_inline_comments():
     assert isinstance(p.tokens[0].tokens[3], sql.Identifier)
 
 
+def test_grouping_identifier_list_with_role():
+    # issue798: ROLE is a non-reserved keyword usable as a column name; in an
+    # identifier list it should be an Identifier like its plain-name peers.
+    p = sqlparse.parse('SELECT a, role, b FROM t')[0]
+    assert isinstance(p.tokens[2], sql.IdentifierList)
+    identifiers = list(p.tokens[2].get_identifiers())
+    assert all(isinstance(i, sql.Identifier) for i in identifiers)
+    assert [i.value for i in identifiers] == ['a', 'role', 'b']
+
+
+def test_grouping_role_keyword_outside_identifier_list():
+    # issue798: ROLE keeps its keyword role outside an identifier list, so e.g.
+    # CREATE ROLE is not misparsed as an identifier aliased to the role name.
+    for sql_str in ('CREATE ROLE myrole', 'SET ROLE admin', 'DROP ROLE r1, r2'):
+        p = sqlparse.parse(sql_str)[0]
+        role = next(t for t in p.flatten() if t.value.upper() == 'ROLE')
+        assert role.ttype is T.Keyword
+        assert not isinstance(role.parent, sql.Identifier)
+
+
 def test_grouping_identifiers_with_operators():
     p = sqlparse.parse('a+b as c from table where (d-e)%2= 1')[0]
     assert len([x for x in p.flatten() if x.ttype == T.Name]) == 5