From 284a0cccd3e1641ab30d347f5212dd7ef991cc00 Mon Sep 17 00:00:00 2001
From: Chinh Nguyen <chinhngt@gmail.com>
Date: Sun, 13 Jan 2019 09:30:05 -0800
Subject: [PATCH] Add fix for pyodbc+mssql (#6621)

* add fix for odbc+mssql

* fix for pylint/pep8
---
 superset/db_engine_specs.py |  7 +++++++
 tests/core_tests.py         | 32 ++++++++++++++++++++++++++++++++
 tests/fixtures/pyodbcRow.py | 12 ++++++++++++
 3 files changed, 51 insertions(+)
 create mode 100644 tests/fixtures/pyodbcRow.py

diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 7cf4a51d9215b..f3726b4d322aa 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -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'
diff --git a/tests/core_tests.py b/tests/core_tests.py
index 886e92fe9feb7..5142ed640026e 100644
--- a/tests/core_tests.py
+++ b/tests/core_tests.py
@@ -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):
@@ -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'
diff --git a/tests/fixtures/pyodbcRow.py b/tests/fixtures/pyodbcRow.py
new file mode 100644
index 0000000000000..f938965d5dc4a
--- /dev/null
+++ b/tests/fixtures/pyodbcRow.py
@@ -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)