[get_df] Adding support for multi-statement SQL (#5086)

superset/models/core.py

@@ -32,6 +32,7 @@
 from sqlalchemy.sql import text
 from sqlalchemy.sql.expression import TextAsFrom
 from sqlalchemy_utils import EncryptedType
+import sqlparse
 
 from superset import app, db, db_engine_specs, security_manager, utils
 from superset.connectors.connector_registry import ConnectorRegistry
@@ -691,8 +692,9 @@ def get_quoter(self):
         return self.get_dialect().identifier_preparer.quote
 
     def get_df(self, sql, schema):
-        sqls = [x.strip() for x in sql.strip().strip(';').split(';')]
+        sqls = [str(s).strip().strip(';') for s in sqlparse.parse(sql)]
         eng = self.get_sqla_engine(schema=schema)
+
         for i in range(len(sqls) - 1):
             eng.execute(sqls[i])
 

tests/model_tests.py

@@ -81,9 +81,15 @@ def test_single_statement(self):
             df = main_db.get_df('SELECT 1', None)
             self.assertEquals(df.iat[0, 0], 1)
 
+            df = main_db.get_df('SELECT 1;', None)
+            self.assertEquals(df.iat[0, 0], 1)
+
     def test_multi_statement(self):
         main_db = self.get_main_database(db.session)
 
         if main_db.backend == 'mysql':
             df = main_db.get_df('USE superset; SELECT 1', None)
             self.assertEquals(df.iat[0, 0], 1)
+
+            df = main_db.get_df("USE superset; SELECT ';';", None)
+            self.assertEquals(df.iat[0, 0], ';')