From 23e823f0c13e34deb4f2d980e4a11b158bdb8bcb Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Mon, 25 Mar 2019 05:44:05 +0200
Subject: [PATCH] Bump SQLAlchemy to 1.3 (#7099)

* Bump sqla to >=1.3.1

* Refine mssql column types to only use N-prefixing when necessary

* make join explicit

* replace set with list

* Add additional test case for N-prefix

* Replace engine with dialect and fix linting error

* Remove unneeded import
---
 docs/installation.rst                         |  5 ---
 requirements.txt                              |  2 +-
 setup.py                                      |  2 +-
 superset/db_engine_specs.py                   | 12 +++++--
 .../4451805bbaa1_remove_double_percents.py    |  4 +--
 tests/db_engine_specs_test.py                 | 36 ++++++++++++++++++-
 6 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/docs/installation.rst b/docs/installation.rst
index 8ef994a23..7bc9f61f3 100644
--- a/docs/installation.rst
+++ b/docs/installation.rst
@@ -415,11 +415,6 @@ You can also use `PyAthena` library(no java required) like this ::
 
 See `PyAthena <https://github.com/laughingman7743/PyAthena#sqlalchemy>`_.
 
-MSSQL
------
-
-Full Unicode support requires SQLAlchemy 1.3 or later.
-
 Snowflake
 ---------
 
diff --git a/requirements.txt b/requirements.txt
index ae4e48751..b20d0c421 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -63,7 +63,7 @@ selenium==3.141.0
 simplejson==3.15.0
 six==1.11.0               # via bleach, cryptography, isodate, pathlib2, polyline, pydruid, python-dateutil, sqlalchemy-utils, wtforms-json
 sqlalchemy-utils==0.32.21
-sqlalchemy==1.2.18
+sqlalchemy==1.3.1
 sqlparse==0.2.4
 unicodecsv==0.14.1
 urllib3==1.22             # via requests, selenium
diff --git a/setup.py b/setup.py
index cc477666f..3c732b8e3 100644
--- a/setup.py
+++ b/setup.py
@@ -100,7 +100,7 @@ setup(
         'retry>=0.9.2',
         'selenium>=3.141.0',
         'simplejson>=3.15.0',
-        'sqlalchemy>=1.2.18, <1.3.0',
+        'sqlalchemy>=1.3.1,<2.0',
         'sqlalchemy-utils',
         'sqlparse',
         'unicodecsv',
diff --git a/superset/db_engine_specs.py b/superset/db_engine_specs.py
index 04779049d..8be1a216d 100644
--- a/superset/db_engine_specs.py
+++ b/superset/db_engine_specs.py
@@ -46,7 +46,7 @@ from sqlalchemy.engine import create_engine
 from sqlalchemy.engine.url import make_url
 from sqlalchemy.sql import quoted_name, text
 from sqlalchemy.sql.expression import TextAsFrom
-from sqlalchemy.types import UnicodeText
+from sqlalchemy.types import String, UnicodeText
 import sqlparse
 from werkzeug.utils import secure_filename
 
@@ -1423,10 +1423,16 @@ class MssqlEngineSpec(BaseEngineSpec):
             data = [[elem for elem in r] for r in data]
         return data
 
+    column_types = [
+        (String(), re.compile(r'^(?<!N)((VAR){0,1}CHAR|TEXT|STRING)', re.IGNORECASE)),
+        (UnicodeText(), re.compile(r'^N((VAR){0,1}CHAR|TEXT)', re.IGNORECASE)),
+    ]
+
     @classmethod
     def get_sqla_column_type(cls, type_):
-        if isinstance(type_, str) and re.match(r'^N(VAR){0-1}CHAR', type_):
-            return UnicodeText()
+        for sqla_type, regex in cls.column_types:
+            if regex.match(type_):
+                return sqla_type
         return None
 
 
diff --git a/superset/migrations/versions/4451805bbaa1_remove_double_percents.py b/superset/migrations/versions/4451805bbaa1_remove_double_percents.py
index c38f0478b..d1edcda4a 100644
--- a/superset/migrations/versions/4451805bbaa1_remove_double_percents.py
+++ b/superset/migrations/versions/4451805bbaa1_remove_double_percents.py
@@ -66,8 +66,8 @@ def replace(source, target):
 
     query = (
         session.query(Slice, Database)
-        .join(Table)
-        .join(Database)
+        .join(Table, Slice.datasource_id == Table.id)
+        .join(Database, Table.database_id == Database.id)
         .filter(Slice.datasource_type == 'table')
         .all()
     )
diff --git a/tests/db_engine_specs_test.py b/tests/db_engine_specs_test.py
index 1390f5092..b5d591e52 100644
--- a/tests/db_engine_specs_test.py
+++ b/tests/db_engine_specs_test.py
@@ -17,7 +17,9 @@
 import inspect
 
 import mock
-from sqlalchemy import column
+from sqlalchemy import column, select, table
+from sqlalchemy.dialects.mssql import pymssql
+from sqlalchemy.types import String, UnicodeText
 
 from superset import db_engine_specs
 from superset.db_engine_specs import (
@@ -346,3 +348,35 @@ class DbEngineSpecsTestCase(SupersetTestCase):
         self.assertEqual(label.quote, True)
         label_expected = '3b26974078683be078219674eeb8f5'
         self.assertEqual(label, label_expected)
+
+    def test_mssql_column_types(self):
+        def assert_type(type_string, type_expected):
+            type_assigned = MssqlEngineSpec.get_sqla_column_type(type_string)
+            if type_expected is None:
+                self.assertIsNone(type_assigned)
+            else:
+                self.assertIsInstance(type_assigned, type_expected)
+
+        assert_type('INT', None)
+        assert_type('STRING', String)
+        assert_type('CHAR(10)', String)
+        assert_type('VARCHAR(10)', String)
+        assert_type('TEXT', String)
+        assert_type('NCHAR(10)', UnicodeText)
+        assert_type('NVARCHAR(10)', UnicodeText)
+        assert_type('NTEXT', UnicodeText)
+
+    def test_mssql_where_clause_n_prefix(self):
+        dialect = pymssql.dialect()
+        spec = MssqlEngineSpec
+        str_col = column('col', type_=spec.get_sqla_column_type('VARCHAR(10)'))
+        unicode_col = column('unicode_col', type_=spec.get_sqla_column_type('NTEXT'))
+        tbl = table('tbl')
+        sel = select([str_col, unicode_col]).\
+            select_from(tbl).\
+            where(str_col == 'abc').\
+            where(unicode_col == 'abc')
+
+        query = str(sel.compile(dialect=dialect, compile_kwargs={'literal_binds': True}))
+        query_expected = "SELECT col, unicode_col \nFROM tbl \nWHERE col = 'abc' AND unicode_col = N'abc'"  # noqa
+        self.assertEqual(query, query_expected)