How to extract columns name along with table name from complex query having subquery

[cehsonu100]

How to extract columns name from complex query having subquery like below:

SELECT "Gender" AS "Gender", "Ethnicity" AS "Ethnicity", count(*) AS "count" FROM (Select distinct avg(t3."Age") As "Age", t1."Ethnicity" As "Ethnicity", t2."Gender" As "Gender" from (SELECT ethnicityid, genderid, age As "Age" FROM schema1.t3 FINAL where (__USER_TABLE_SECURITY_t3_USER_TABLE_SECURITY__) ) as t3 inner join (SELECT name As "Ethnicity", ethnicityid FROM schema1.t1 FINAL where (t1.code = 'AP' and __USER_TABLE_SECURITY_t1_USER_TABLE_SECURITY__) ) as t1 on t3.ethnicityid = t1.ethnicityid inner join (SELECT name As "Gender", genderid FROM schema1.t2 FINAL where (__USER_TABLE_SECURITY_t2_USER_TABLE_SECURITY__) ) as t2 on t3.genderid = t2.genderid where (1 = 1) group by "Ethnicity", "Gender") AS "virtual_table" GROUP BY "Gender", "Ethnicity" ORDER BY "count" DESC LIMIT 1000;

[ChrisJaunes]

You can use "splparse.parse" to parse this complex query, then you can iterate over tokens and get cloumns from "select token" and "from token". Note that it is handled according to different cloumn types.

[ChrisJaunes]

more, you can refer to https://github.com/andialbrecht/sqlparse/blob/master/examples/extract_table_names.py

[ChrisJaunes]

If you need to extract the outermost columns, some modifications are required

for example:

if __name__ == "__main__":
    sql = """SELECT "Gender" AS "Gender", "Ethnicity" AS "Ethnicity", count(*) AS "count" FROM (Select distinct avg(t3."Age") As "Age", t1."Ethnicity" As "Ethnicity", t2."Gender" As "Gender" from (SELECT ethnicityid, genderid, age As "Age" FROM schema1.t3 FINAL where (__USER_TABLE_SECURITY_t3_USER_TABLE_SECURITY__) ) as t3 inner join (SELECT name As "Ethnicity", ethnicityid FROM schema1.t1 FINAL where (t1.code = 'AP' and __USER_TABLE_SECURITY_t1_USER_TABLE_SECURITY__) ) as t1 on t3.ethnicityid = t1.ethnicityid inner join (SELECT name As "Gender", genderid FROM schema1.t2 FINAL where (__USER_TABLE_SECURITY_t2_USER_TABLE_SECURITY__) ) as t2 on t3.genderid = t2.genderid where (1 = 1) group by "Ethnicity", "Gender") AS "virtual_table" GROUP BY "Gender", "Ethnicity" ORDER BY "count" DESC LIMIT 1000;"""
    import sqlparse
    statement = sqlparse.parse(sql)[0]
    raw_columns = []
    USE_SELECT, USE_FROM = False, False
    for token in statement:
        if token.ttype == sqlparse.tokens.Keyword and token.value == "FROM":
            USE_FROM = True
            break
        if USE_SELECT and not (token.is_whitespace or token.match(sqlparse.tokens.Punctuation, ',')):
            raw_columns.append(token)
        if token.ttype == sqlparse.tokens.DML and token.value == "SELECT":
            USE_SELECT = True

    def extract_table_identifiers(token_stream):
        for item in token_stream:
            if isinstance(item, sqlparse.sql.IdentifierList):
                for identifier in item.get_identifiers():
                    yield identifier.get_name()
            elif isinstance(item, sqlparse.sql.Identifier):
                yield item.get_name()
            elif item.ttype is sqlparse.tokens.Keyword:
                yield item.value

    print(list(extract_table_identifiers(raw_columns)))

output

['Gender', 'Ethnicity', 'count']

[Rajdeep-689]

Hi @ChrisJaunes . I already have the tokenized form of a sql query in json struture. I want to extract the table and column name it's not coming. the below I have written the code can you please check why the value of current_table & column_name is not flowing to the next elif statements.

def traverse(node, current_table=None):
"""
Traverse the SQL structure to extract tables and columns.
"""
if isinstance(node, dict):
for key, value in node.items():
# If it's a table reference, set current_table
if key == 'table_reference' and isinstance(value, list):
# Join parts like table alias and table name
table_name_parts = [v.get('naked_identifier') for v in value if v.get('naked_identifier')]
table_name = ".".join(table_name_parts)
if table_name:
current_table = table_name
print(f"Table detected: {current_table}") # Debugging

        # Handle column references
        elif key == 'column_reference' and isinstance(value, dict):
            # Extract column name directly from the structure
            column_name = value.get('naked_identifier')  # Assuming column name is stored here
            if column_name:  # Ensure column name is valid
                print(f"Column detected: {column_name} for table: {current_table}")  # Debugging
                if current_table:
                    if current_table not in tables_columns_map:
                        tables_columns_map[current_table] = []
                    tables_columns_map[current_table].append(column_name)
                else:
                    print(f"Error: current_table is None when processing column {column_name}")

        # Recurse into subnodes
        elif isinstance(value, (list, dict)):
            traverse(value, current_table)  # Continue processing

elif isinstance(node, list):
    for item in node:
        traverse(item, current_table)

Example usage:

parsed_data = {
"select_statement": [
{"column_reference": {"naked_identifier": "MI_KEY_ACC_SWITCHING"}},
{"column_reference": {"naked_identifier": "CURRENCY"}},
{"table_reference": [{"naked_identifier": "A_MPAY_BUKDB.ACC_SWITCHING"}]},
{"column_reference": {"naked_identifier": "DEPARTMENT"}}
]
}

tables_columns_map = {}
traverse(parsed_data)
print("Final tables_columns_map:")
print(tables_columns_map)

[Rajdeep-689]

Also I tried with raw sql query,

sql_query = """
INSERT INTO A_MPAY_BUKDB.ACC_SWITCHING
(MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT)
SELECT MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT
FROM U_MPAY_BUKDB.ACC_SWITCHING
WHERE EFF_STA_DTE = (
SELECT CURR_BUSS_DATE
FROM A_MPAY_BUKDB.MI_DATE_CONTROL
WHERE SUITE_ID = 'MPY'
);

UPDATE A_MPAY_BUKDB.ACC_SWITCHING
SET CURRENCY = 'USD', DEPARTMENT = 'HR'
FROM U_MPAY_BUKDB.ACC_SWITCHING
WHERE EFF_STA_DTE = '2023-01-01';
"""

So basically the output should come like the below: Can we get this with any way???

Table: A_MPAY_BUKDB.ACC_SWITCHING
Columns from main query:

MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT, EFF_STA_DTE

Columns from subqueries:
CURR_BUSS_DATE, SUITE_ID

---

Table: U_MPAY_BUKDB.ACC_SWITCHING
Columns from main query:

MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT, EFF_STA_DTE

Columns from subqueries:

Table: A_MPAY_BUKDB.MI_DATE_CONTROL
Columns from main query:

Columns from subqueries:
CURR_BUSS_DATE, SUITE_ID

[ChrisJaunes]

Hi @Rajdeep-689
Your code is missing indentation.

It should be fine to get the first non-blank tokens after FROM

def tmp(parsed):
    select_seen = False
    for item in parsed.tokens:
        if select_seen and item.ttype != sqlparse.tokens.Text.Whitespace:
            return item.value
        if item.ttype is sqlparse.tokens.Keyword and item.value.upper() == 'FROM':
            select_seen = True
    return "ERR"

tmp(sqlparse.parse("""
INSERT INTO A_MPAY_BUKDB.ACC_SWITCHING
(MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT)
SELECT MI_KEY_ACC_SWITCHING, CURRENCY, DEPARTMENT
FROM U_MPAY_BUKDB.ACC_SWITCHING
WHERE EFF_STA_DTE = (
SELECT CURR_BUSS_DATE
FROM A_MPAY_BUKDB.MI_DATE_CONTROL
WHERE SUITE_ID = 'MPY'
);

UPDATE A_MPAY_BUKDB.ACC_SWITCHING
SET CURRENCY = 'USD', DEPARTMENT = 'HR'
FROM U_MPAY_BUKDB.ACC_SWITCHING
WHERE EFF_STA_DTE = '2023-01-01';
""")[0])

[stegarth]

Posting this here should it benefit someone else in the future.

When using the method proposed in this pull request. I kept encountering the error:

AttributeError: 'Token' object has no attribute 'get_name'

when attempting to extract columns from a query of the form:

SELECT id, segment FROM db.schema.table

Evidently segment is treated as a Keyword thus causing my issue.

To get around this I modified @ChrisJaunes's code to the include the following:

import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Keyword, DML

def extract_select_part(parsed):
    select_seen = False
    for item in parsed.tokens:
        if item.ttype is Keyword and item.value.upper() == 'FROM':
            return
        if select_seen:
            yield item
        if item.ttype is DML and item.value.upper() == 'SELECT':
            select_seen = True

def col_name(item):
    if isinstance(item, IdentifierList):
        for id_ in item.get_identifiers():
            yield from col_name(id_)
    elif isinstance(item, Identifier):
        yield item.get_name()
    elif item.ttype is Keyword:
        yield item.value

def extract_column_identifiers(token_stream):
    for item in token_stream:
        yield from col_name(item)

def extract_columns(sql):
    stream = extract_select_part(sqlparse.parse(sql)[0])
    return list(extract_column_identifiers(stream))

which resolved the issue for me. Not sure if this is the best way to resolve that problem but it seems to be working with whatever I throw at it 🤷