From 1ee3c57e48ada3a82cf96077f0788b6ee6674d03 Mon Sep 17 00:00:00 2001
From: Timi Fasubaa <timi.fasubaa@airbnb.com>
Date: Tue, 28 Aug 2018 20:23:35 -0700
Subject: [PATCH] fix sqlparse bug and refactor

(cherry picked from commit 1e94215fd646d7af84b87921fc2b0b0103eadc82)
---
 superset/sql_parse.py    |  4 ++--
 tests/sql_parse_tests.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 54320d6c18e68..d106f38ac7a5b 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -11,7 +11,7 @@
 from sqlparse.sql import Identifier, IdentifierList
 from sqlparse.tokens import Keyword, Name
 
-RESULT_OPERATIONS = {'UNION', 'INTERSECT', 'EXCEPT'}
+RESULT_OPERATIONS = {'UNION', 'INTERSECT', 'EXCEPT', 'SELECT'}
 ON_KEYWORD = 'ON'
 PRECEDES_TABLE_NAME = {'FROM', 'JOIN', 'DESC', 'DESCRIBE', 'WITH'}
 
@@ -125,7 +125,7 @@ def __extract_from_token(self, token):
             if not table_name_preceding_token:
                 continue
 
-            if item.ttype in Keyword:
+            if item.ttype in Keyword or item.value == ',':
                 if (self.__is_result_operation(item.value) or
                         item.value.upper() == ON_KEYWORD):
                     table_name_preceding_token = False
diff --git a/tests/sql_parse_tests.py b/tests/sql_parse_tests.py
index 1ef08a6333211..6f9290fc4eeb3 100644
--- a/tests/sql_parse_tests.py
+++ b/tests/sql_parse_tests.py
@@ -326,3 +326,40 @@ def test_complex_extract_tables2(self):
         self.assertEquals(
             {'table_a', 'table_b', 'table_c'},
             self.extract_tables(query))
+
+    def test_complex_extract_tables3(self):
+        query = """SELECT somecol AS somecol
+            FROM
+              (WITH bla AS
+                 (SELECT col_a
+                  FROM a
+                  WHERE 1=1
+                    AND column_of_choice NOT IN
+                      ( SELECT interesting_col
+                       FROM b ) ),
+                    rb AS
+                 ( SELECT yet_another_column
+                  FROM
+                    ( SELECT a
+                     FROM c
+                     GROUP BY the_other_col ) not_table
+                  LEFT JOIN bla foo ON foo.prop = not_table.bad_col0
+                  WHERE 1=1
+                  GROUP BY not_table.bad_col1 ,
+                           not_table.bad_col2 ,
+                  ORDER BY not_table.bad_col_3 DESC , not_table.bad_col4 ,
+                                                    not_table.bad_col5) SELECT random_col
+               FROM d
+               WHERE 1=1
+               UNION ALL SELECT even_more_cols
+               FROM e
+               WHERE 1=1
+               UNION ALL SELECT lets_go_deeper
+               FROM f
+               WHERE 1=1
+            WHERE 2=2
+            GROUP BY last_col
+            LIMIT 50000;"""
+        self.assertEquals(
+            {'a', 'b', 'c', 'd', 'e', 'f'},
+            self.extract_tables(query))