python sqlparse获取血缘 代码发我

以下是使用sqlparse获取SQL语句血缘的示例代码：

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

# 解析SQL语句，获取表名和列名
def extract_table_and_column(sql):
    tables = set()
    columns = set()
    parsed = sqlparse.parse(sql)
    for stmt in parsed:
        if stmt.token_first(skip_ws=True).ttype is DML:
            for token in stmt.tokens:
                if token.ttype is Keyword:
                    if token.value.upper() == 'FROM' or token.value.upper() == 'JOIN':
                        tables.add(token.next_sibling.next_sibling.value.strip())
                    elif token.value.upper() == 'SELECT':
                        for item in token.parent.tokens:
                            if item.ttype is IdentifierList:
                                for identifier in item.get_identifiers():
                                    columns.add(identifier.value.strip())
                            elif isinstance(item, Identifier) and item.ttype is not None:
                                columns.add(item.value.strip())
    return tables, columns

# 获取SQL语句血缘
def get_sql_lineage(sql):
    tables, columns = extract_table_and_column(sql)
    lineage = {}
    for table in tables:
        lineage[table] = set()
        for column in columns:
            if '.' in column and column.split('.')[0] == table:
                lineage[table].add(column.split('.')[1])
    return lineage

# 示例SQL语句
sql = "SELECT a.name, b.age FROM table1 a JOIN table2 b ON a.id = b.id WHERE a.id = 1"

# 获取SQL语句血缘
lineage = get_sql_lineage(sql)
print(lineage)

这段代码会解析输入的SQL语句，并提取出其中涉及的表名和列名，最终返回一个字典，其中键为表名，值为该表所涉及的列名集合。