import sqlparse
from sqlparse.sql import IdentifierList, Identifier
from sqlparse.tokens import Keyword, DML
 
 
# 支持的join方式
ALL_JOIN_TYPE = (‘LEFT JOIN‘, ‘RIGHT JOIN‘, ‘INNER JOIN‘, ‘FULL JOIN‘, ‘LEFT OUTER JOIN‘, ‘FULL OUTER JOIN‘)
 
 
def is_subselect(parsed):
    """
    是否子查询
    :param parsed:
    :return:
    """
    if not parsed.is_group:
        return False
    for item in parsed.tokens:
        if item.ttype is DML and item.value.upper() == ‘SELECT‘:
            return True
    return False
 
 
def extract_from_part(parsed):
    """
    提取from之后模块
    :param parsed:
    :return:
    """
    from_seen = False
    for item in parsed.tokens:
        if from_seen:
            if is_subselect(item):
                for x in extract_from_part(item):
                    yield x
            elif item.ttype is Keyword:
                from_seen = False
                continue
            else:
                yield item
        elif item.ttype is Keyword and item.value.upper() == ‘FROM‘:
            from_seen = True
 
 
def extract_join_part(parsed):
    """
    提交join之后模块
    :param parsed:
    :return:
    """
    flag = False
    for item in parsed.tokens:
        if flag:
            if item.ttype is Keyword:
                flag = False
                continue
            else:
                yield item
        if item.ttype is Keyword and item.value.upper() in ALL_JOIN_TYPE:
            flag = True
 
 
def extract_table_identifiers(token_stream):
    for item in token_stream:
        if isinstance(item, IdentifierList):
            for identifier in item.get_identifiers():
                yield identifier.get_name()
        elif isinstance(item, Identifier):
            yield item.get_name()
        elif item.ttype is Keyword:
            yield item.value
 
 
def extract_tables(sql):
    """
    提取sql中的表名（select语句）
    :param sql:
    :return:
    """
    from_stream = extract_from_part(sqlparse.parse(sql)[0])
    join_stream = extract_join_part(sqlparse.parse(sql)[0])
    return list(extract_table_identifiers(from_stream)) + list(extract_table_identifiers(join_stream))
 
 
 
print(extract_tables("select x1, x2 from a left join b on a.id = b.id right join c on c.id = a.id"))
print(extract_tables("select x1, x2 from (select x1, x2 from (select x1, x2 from a)) left join b on a.id=b.id"))