Skip to content

Commit

Permalink
Add fix for pyodbc+mssql (#6621)
Browse files Browse the repository at this point in the history
* add fix for odbc+mssql

* fix for pylint/pep8
  • Loading branch information
chinhngt authored and mistercrunch committed Jan 13, 2019
1 parent ae6217b commit 284a0cc
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 0 deletions.
7 changes: 7 additions & 0 deletions superset/db_engine_specs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1287,6 +1287,13 @@ class MssqlEngineSpec(BaseEngineSpec):
def convert_dttm(cls, target_type, dttm):
return "CONVERT(DATETIME, '{}', 126)".format(dttm.isoformat())

@classmethod
def fetch_data(cls, cursor, limit):
data = super(MssqlEngineSpec, cls).fetch_data(cursor, limit)
if len(data) != 0 and type(data[0]).__name__ == 'Row':
data = [[elem for elem in r] for r in data]
return data


class AthenaEngineSpec(BaseEngineSpec):
engine = 'awsathena'
Expand Down
32 changes: 32 additions & 0 deletions tests/core_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,14 @@
from superset import dataframe, db, jinja_context, security_manager, sql_lab
from superset.connectors.sqla.models import SqlaTable
from superset.db_engine_specs import BaseEngineSpec
from superset.db_engine_specs import MssqlEngineSpec
from superset.models import core as models
from superset.models.sql_lab import Query
from superset.utils import core as utils
from superset.utils.core import get_main_database
from superset.views.core import DatabaseView
from .base_tests import SupersetTestCase
from .fixtures.pyodbcRow import Row


class CoreTests(SupersetTestCase):
Expand Down Expand Up @@ -673,6 +675,36 @@ def test_dataframe_timezone(self):
{'data': pd.Timestamp('2017-11-18 22:06:30.061810+0100', tz=tz)},
)

def test_mssql_engine_spec_pymssql(self):
# Test for case when tuple is returned (pymssql)
data = [(1, 1, datetime.datetime(2017, 10, 19, 23, 39, 16, 660000)),
(2, 2, datetime.datetime(2018, 10, 19, 23, 39, 16, 660000))]
df = dataframe.SupersetDataFrame(
list(data),
[['col1'], ['col2'], ['col3']],
MssqlEngineSpec)
data = df.data
self.assertEqual(len(data), 2)
self.assertEqual(data[0],
{'col1': 1,
'col2': 1,
'col3': pd.Timestamp('2017-10-19 23:39:16.660000')})

def test_mssql_engine_spec_odbc(self):
# Test for case when pyodbc.Row is returned (msodbc driver)
data = [Row((1, 1, datetime.datetime(2017, 10, 19, 23, 39, 16, 660000))),
Row((2, 2, datetime.datetime(2018, 10, 19, 23, 39, 16, 660000)))]
df = dataframe.SupersetDataFrame(
list(data),
[['col1'], ['col2'], ['col3']],
MssqlEngineSpec)
data = df.data
self.assertEqual(len(data), 2)
self.assertEqual(data[0],
{'col1': 1,
'col2': 1,
'col3': pd.Timestamp('2017-10-19 23:39:16.660000')})

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'
Expand Down
12 changes: 12 additions & 0 deletions tests/fixtures/pyodbcRow.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# pylint: disable=C,R,W


class Row(object):
def __init__(self, values):
self.values = values

def __name__(self):
return 'Row'

def __iter__(self):
return (item for item in self.values)

0 comments on commit 284a0cc

Please sign in to comment.