diff --git a/superset/models/core.py b/superset/models/core.py index a9b9baf0d352f..ebce5fcb3103c 100644 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -691,9 +691,13 @@ def get_quoter(self): return self.get_dialect().identifier_preparer.quote def get_df(self, sql, schema): - sql = sql.strip().strip(';') + sqls = [str(s).strip().strip(';') for s in sqlparse.parse(sql)] eng = self.get_sqla_engine(schema=schema) - df = pd.read_sql_query(sql, eng) + + for i in range(len(sqls) - 1): + eng.execute(sqls[i]) + + df = pd.read_sql_query(sqls[-1], eng) def needs_conversion(df_series): if df_series.empty: diff --git a/tests/model_tests.py b/tests/model_tests.py index b223aa7e556ee..c0fe0bd774931 100644 --- a/tests/model_tests.py +++ b/tests/model_tests.py @@ -110,6 +110,9 @@ def test_single_statement(self): main_db = self.get_main_database(db.session) if main_db.backend == 'mysql': + 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) @@ -120,6 +123,13 @@ def test_multi_statement(self): df = main_db.get_df("USE superset; SELECT ';';", None) self.assertEquals(df.iat[0, 0], ';') + 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], ';') + + class SqlaTableModelTestCase(SupersetTestCase): def test_get_timestamp_expression(self):