How to parse any SQL get columns names and table name using SQL parser in python3

Question

I am able to get the column names and table name from using sql parse for only simple select SQL's.

Can somebody help how can get the column names and table name from any complex SQL's.

Answer 1

Here is a solution for extracting column names from complex sql select statements. Python 3.9

import sqlparse

def get_query_columns(sql):
    stmt = sqlparse.parse(sql)[0]
    columns = []
    column_identifiers = []

    # get column_identifieres
    in_select = False
    for token in stmt.tokens:
        if isinstance(token, sqlparse.sql.Comment):
            continue
        if str(token).lower() == 'select':
            in_select = True
        elif in_select and token.ttype is None:
            for identifier in token.get_identifiers():
                column_identifiers.append(identifier)
            break

    # get column names
    for column_identifier in column_identifiers:
        columns.append(column_identifier.get_name())

    return columns

def test():
    sql = '''
select
   a.a,
   replace(coalesce(a.b, 'x'), 'x', 'y') as jim,
   a.bla as sally  -- some comment
from
   table_a as a
where
   c > 20
'''
    print(get_query_columns(sql))

test()
# outputs: ['a', 'jim', 'sally']