Remove comments from queries in SQL Lab that break Explore view (apac…

…he#4413) * Remove comments from queries in SQL Lab that break Explore view This fixes an issue where comments on the last line of the source query comment out the closing parenthesis of the subquery. * Add test case for SqlaTable with query with comment This test ensures that comments in the query are removed when calling SqlaTable.get_from_clause(). * Add missing blank line class definition (PEP8)
wenchma · Feb 19, 2018 · d2d135b · d2d135b
1 parent 7931893
commit d2d135b
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
@@ -428,6 +428,7 @@ def get_from_clause(self, template_processor=None, db_engine_spec=None):
                 from_sql = template_processor.process_template(from_sql)
             if db_engine_spec:
                 from_sql = db_engine_spec.escape_sql(from_sql)
+            from_sql = sqlparse.format(from_sql, strip_comments=True)
             return TextAsFrom(sa.text(from_sql), []).alias('expr_qry')
         return self.get_sqla_table()
 

diff --git a/tests/core_tests.py b/tests/core_tests.py
@@ -18,6 +18,7 @@
 from flask import escape
 import pandas as pd
 import psycopg2
+from six import text_type
 import sqlalchemy as sqla
 
 from superset import appbuilder, dataframe, db, jinja_context, sm, sql_lab, utils
@@ -870,6 +871,13 @@ def test_dataframe_timezone(self):
             {'data': pd.Timestamp('2017-11-18 22:06:30.061810+0100', tz=tz)},
         )
 
+    def test_comments_in_sqlatable_query(self):
+        clean_query = "SELECT '/* val 1 */' as c1, '-- val 2' as c2 FROM tbl"
+        commented_query = '/* comment 1 */' + clean_query + '-- comment 2'
+        table = SqlaTable(sql=commented_query)
+        rendered_query = text_type(table.get_from_clause())
+        self.assertEqual(clean_query, rendered_query)
+
 
 if __name__ == '__main__':
     unittest.main()