Fix SQLParse when query has reserved word (#7305)

* Fix SQLParse when query has reserved word

* Fix indent

superset/sql_parse.py

@@ -86,6 +86,9 @@ def __is_identifier(token):
         return isinstance(token, (IdentifierList, Identifier))
 
     def __process_identifier(self, identifier):
+        if not self.__is_identifier(identifier):
+            identifier = Identifier([identifier])
+
         # exclude subselects
         if '(' not in str(identifier):
             table_name = self.__get_full_name(identifier)

tests/sql_parse_tests.py

@@ -462,3 +462,28 @@ def test_messy_breakdown_statements(self):
             'SELECT * FROM ab_user LIMIT 1',
         ]
         self.assertEquals(statements, expected)
+
+    def test_extract_from_statement(self):
+        query = 'SELECT a, b FROM c WHERE d'
+        result = sql_parse.ParsedQuery(query)._table_names
+        expected = {'c'}
+        self.assertEquals(result, expected)
+
+    def test_extract_from_statement_with_reserved(self):
+        query = """
+        WITH
+            columns AS (SELECT metric FROM a),
+            rows AS (SELECT metric FROM b)
+        SELECT
+            c.metric AS m1,
+            r.metric AS m2
+        FROM columns c
+        JOIN rows r
+        """
+        result = sql_parse.ParsedQuery(query)._table_names
+        expected = {'a', 'b'}
+        self.assertEquals(result, expected)
+
+
+if __name__ == '__main__':
+    unittest.main()