diff --git a/superset/sql_parse.py b/superset/sql_parse.py index 790371ae35706..54320d6c18e68 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -12,6 +12,7 @@ from sqlparse.tokens import Keyword, Name RESULT_OPERATIONS = {'UNION', 'INTERSECT', 'EXCEPT'} +ON_KEYWORD = 'ON' PRECEDES_TABLE_NAME = {'FROM', 'JOIN', 'DESC', 'DESCRIBE', 'WITH'} @@ -125,7 +126,8 @@ def __extract_from_token(self, token): continue if item.ttype in Keyword: - if SupersetQuery.__is_result_operation(item.value): + if (self.__is_result_operation(item.value) or + item.value.upper() == ON_KEYWORD): table_name_preceding_token = False continue # FROM clause is over diff --git a/tests/sql_parse_tests.py b/tests/sql_parse_tests.py index c9368bb362a27..1ef08a6333211 100644 --- a/tests/sql_parse_tests.py +++ b/tests/sql_parse_tests.py @@ -298,3 +298,31 @@ def multistatement(self): query = 'SELECT * FROM t1; SELECT * FROM t2;' self.assertEquals({'t1', 't2'}, self.extract_tables(query)) + + def test_complex_extract_tables(self): + query = """SELECT sum(m_examples) AS "sum__m_example" + FROM + (SELECT COUNT(DISTINCT id_userid) AS m_examples, + some_more_info + FROM my_b_table b + JOIN my_t_table t ON b.ds=t.ds + JOIN my_l_table l ON b.uid=l.uid + WHERE b.rid IN + (SELECT other_col + FROM inner_table) + AND l.bla IN ('x', 'y') + GROUP BY 2 + ORDER BY 2 ASC) AS "meh" + ORDER BY "sum__m_example" DESC + LIMIT 10;""" + self.assertEquals( + {'my_l_table', 'my_b_table', 'my_t_table', 'inner_table'}, + self.extract_tables(query)) + + def test_complex_extract_tables2(self): + query = """SELECT * + FROM table_a AS a, table_b AS b, table_c as c + WHERE a.id = b.id and b.id = c.id""" + self.assertEquals( + {'table_a', 'table_b', 'table_c'}, + self.extract_tables(query))