pt
/
superset
mirror of https://github.com/apache/superset.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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

This commit is contained in:
John Bodley 2018-05-23 11:40:25 -07:00 committed by GitHub
parent 4c44223234
commit 17d6464aa9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 20 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
superset/models/core.py
View File

@ -690,9 +690,13 @@ class Database(Model, AuditMixinNullable, ImportMixin):
return self.get_dialect().identifier_preparer.quote
def get_df(self, sql, schema):
sql = sql.strip().strip(';')
sqls = [x.strip() for x in sql.strip().strip(';').split(';')]
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:

14
tests/model_tests.py
View File

@ -106,6 +106,20 @@ class DatabaseModelTestCase(SupersetTestCase):
self.assertEquals(d.get('P1D').function, 'DATE({col})')
self.assertEquals(d.get('Time Column').function, '{col}')
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)
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)
class SqlaTableModelTestCase(SupersetTestCase):