mirror of
https://github.com/apache/superset.git
synced 2024-09-17 11:09:47 -04:00
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
|
#!/usr/bin/env python
|
||
|
# -*- coding: utf-8 -*-
|
||
|
#
|
||
|
# Copyright (C) 2016 Andi Albrecht, albrecht.andi@gmail.com
|
||
|
#
|
||
|
# This example is part of python-sqlparse and is released under
|
||
|
# the BSD License: http://www.opensource.org/licenses/bsd-license.php
|
||
|
#
|
||
|
# This example illustrates how to extract table names from nested
|
||
|
# SELECT statements.
|
||
|
#
|
||
|
# See:
|
||
|
# http://groups.google.com/group/sqlparse/browse_thread/thread/b0bd9a022e9d4895
|
||
|
|
||
|
import sqlparse
|
||
|
from sqlparse.sql import IdentifierList, Identifier
|
||
|
from sqlparse.tokens import Keyword, DML
|
||
|
|
||
|
|
||
|
def is_subselect(parsed):
|
||
|
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_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:
|
||
|
raise StopIteration
|
||
|
else:
|
||
|
yield item
|
||
|
elif item.ttype is Keyword and item.value.upper() == 'FROM':
|
||
|
from_seen = 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()
|
||
|
# It's a bug to check for Keyword here, but in the example
|
||
|
# above some tables names are identified as keywords...
|
||
|
elif item.ttype is Keyword:
|
||
|
yield item.value
|
||
|
|
||
|
|
||
|
# TODO(bkyryliuk): add logic to support joins and unions.
|
||
|
def extract_tables(sql):
|
||
|
stream = extract_from_part(sqlparse.parse(sql)[0])
|
||
|
return list(extract_table_identifiers(stream))
|