From 59cb96dda2124b6da2634489edb06d2418e95304 Mon Sep 17 00:00:00 2001 From: John Bodley <4567245+john-bodley@users.noreply.github.com> Date: Tue, 29 May 2018 14:20:17 -0700 Subject: [PATCH] [get_df] Adding support for multi-statement SQL (#5086) (cherry picked from commit 0511d1f38d99d98bcc0d179e0578d20f4d4c2985) --- superset/models/core.py | 4 +++- tests/model_tests.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/superset/models/core.py b/superset/models/core.py index 431d5366d23d7..ff03b6f998d30 100644 --- a/superset/models/core.py +++ b/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]) diff --git a/tests/model_tests.py b/tests/model_tests.py index f25319556ea32..de6d0bf8c4ca1 100644 --- a/tests/model_tests.py +++ b/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], ';')