Merge "Add JSON support for mssql"

10 files changed, 310 insertions, 34 deletions

diff --git a/lib/sqlalchemy/dialects/mssql/__init__.py b/lib/sqlalchemy/dialects/mssql/__init__.py
index 283c92eca..d987efa51 100644
--- a/lib/sqlalchemy/dialects/mssql/__init__.py
+++ b/lib/sqlalchemy/dialects/mssql/__init__.py
@@ -21,6 +21,7 @@ from .base import DECIMAL
 from .base import FLOAT
 from .base import IMAGE
 from .base import INTEGER
+from .base import JSON
 from .base import MONEY
 from .base import NCHAR
 from .base import NTEXT
@@ -47,6 +48,7 @@ base.dialect = dialect = pyodbc.dialect
 
 
 __all__ = (
+    "JSON",
     "INTEGER",
     "BIGINT",
     "SMALLINT",
diff --git a/lib/sqlalchemy/dialects/mssql/base.py b/lib/sqlalchemy/dialects/mssql/base.py
index 0ec6cf8a3..ab6e19cf4 100644
--- a/lib/sqlalchemy/dialects/mssql/base.py
+++ b/lib/sqlalchemy/dialects/mssql/base.py
@@ -716,6 +716,9 @@ import operator
 import re
 
 from . import information_schema as ischema
+from .json import JSON
+from .json import JSONIndexType
+from .json import JSONPathType
 from ... import exc
 from ... import schema as sa_schema
 from ... import Sequence
@@ -1453,6 +1456,11 @@ class MSTypeCompiler(compiler.GenericTypeCompiler):
     def visit_BIT(self, type_, **kw):
         return "BIT"
 
+    def visit_JSON(self, type_, **kw):
+        # this is a bit of a break with SQLAlchemy's convention of
+        # "UPPERCASE name goes to UPPERCASE type name with no modification"
+        return self._extend("NVARCHAR", type_, length="max")
+
     def visit_MONEY(self, type_, **kw):
         return "MONEY"
 
@@ -2010,6 +2018,65 @@ class MSSQLCompiler(compiler.SQLCompiler):
             self.process(binary.right),
         )
 
+    def _render_json_extract_from_binary(self, binary, operator, **kw):
+        # note we are intentionally calling upon the process() calls in the
+        # order in which they appear in the SQL String as this is used
+        # by positional parameter rendering
+
+        if binary.type._type_affinity is sqltypes.JSON:
+            return "JSON_QUERY(%s, %s)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+
+        # as with other dialects, start with an explicit test for NULL
+        case_expression = "CASE JSON_VALUE(%s, %s) WHEN NULL THEN NULL" % (
+            self.process(binary.left, **kw),
+            self.process(binary.right, **kw),
+        )
+
+        if binary.type._type_affinity is sqltypes.Integer:
+            type_expression = "ELSE CAST(JSON_VALUE(%s, %s) AS INTEGER)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+        elif binary.type._type_affinity is sqltypes.Numeric:
+            type_expression = (
+                "ELSE CAST(JSON_VALUE(%s, %s) AS DECIMAL(10, 6))"
+                % (
+                    self.process(binary.left, **kw),
+                    self.process(binary.right, **kw),
+                )
+            )
+        elif binary.type._type_affinity is sqltypes.Boolean:
+            # the NULL handling is particularly weird with boolean, so
+            # explicitly return numeric (BIT) constants
+            type_expression = (
+                "WHEN 'true' THEN 1 WHEN 'false' THEN 0 ELSE NULL"
+            )
+        elif binary.type._type_affinity is sqltypes.String:
+            # TODO: does this comment (from mysql) apply to here, too?
+            #       this fails with a JSON value that's a four byte unicode
+            #       string.  SQLite has the same problem at the moment
+            type_expression = "ELSE JSON_VALUE(%s, %s)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+        else:
+            # other affinity....this is not expected right now
+            type_expression = "ELSE JSON_QUERY(%s, %s)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+
+        return case_expression + " " + type_expression + " END"
+
+    def visit_json_getitem_op_binary(self, binary, operator, **kw):
+        return self._render_json_extract_from_binary(binary, operator, **kw)
+
+    def visit_json_path_getitem_op_binary(self, binary, operator, **kw):
+        return self._render_json_extract_from_binary(binary, operator, **kw)
+
     def visit_sequence(self, seq, **kw):
         return "NEXT VALUE FOR %s" % self.preparer.format_sequence(seq)
 
@@ -2412,6 +2479,9 @@ class MSDialect(default.DefaultDialect):
     colspecs = {
         sqltypes.DateTime: _MSDateTime,
         sqltypes.Date: _MSDate,
+        sqltypes.JSON: JSON,
+        sqltypes.JSON.JSONIndexType: JSONIndexType,
+        sqltypes.JSON.JSONPathType: JSONPathType,
         sqltypes.Time: TIME,
         sqltypes.Unicode: _MSUnicode,
         sqltypes.UnicodeText: _MSUnicodeText,
@@ -2458,6 +2528,8 @@ class MSDialect(default.DefaultDialect):
         isolation_level=None,
         deprecate_large_types=None,
         legacy_schema_aliasing=False,
+        json_serializer=None,
+        json_deserializer=None,
         **opts
     ):
         self.query_timeout = int(query_timeout or 0)
@@ -2470,6 +2542,8 @@ class MSDialect(default.DefaultDialect):
         super(MSDialect, self).__init__(**opts)
 
         self.isolation_level = isolation_level
+        self._json_serializer = json_serializer
+        self._json_deserializer = json_deserializer
 
     def do_savepoint(self, connection, name):
         # give the DBAPI a push
diff --git a/lib/sqlalchemy/dialects/mssql/json.py b/lib/sqlalchemy/dialects/mssql/json.py
new file mode 100644
index 000000000..2cae3f774
--- /dev/null
+++ b/lib/sqlalchemy/dialects/mssql/json.py
@@ -0,0 +1,125 @@
+from ... import types as sqltypes
+
+# technically, all the dialect-specific datatypes that don't have any special
+# behaviors would be private with names like _MSJson. However, we haven't been
+# doing this for mysql.JSON or sqlite.JSON which both have JSON / JSONIndexType
+# / JSONPathType in their json.py files, so keep consistent with that
+# sub-convention for now.  A future change can update them all to be
+# package-private at once.
+
+
+class JSON(sqltypes.JSON):
+    """MSSQL JSON type.
+
+    MSSQL supports JSON-formatted data as of SQL Server 2016.
+
+    The :class:`_mssql.JSON` datatype at the DDL level will represent the
+    datatype as ``NVARCHAR(max)``, but provides for JSON-level comparison
+    functions as well as Python coercion behavior.
+
+    :class:`_mssql.JSON` is used automatically whenever the base
+    :class:`_types.JSON` datatype is used against a SQL Server backend.
+
+    .. seealso::
+
+        :class:`_types.JSON` - main documenation for the generic
+        cross-platform JSON datatype.
+
+    The :class:`_mssql.JSON` type supports persistence of JSON values
+    as well as the core index operations provided by :class:`_types.JSON`
+    datatype, by adapting the operations to render the ``JSON_VALUE``
+    or ``JSON_QUERY`` functions at the database level.
+
+    The SQL Server :class:`_mssql.JSON` type necessarily makes use of the
+    ``JSON_QUERY`` and ``JSON_VALUE`` functions when querying for elements
+    of a JSON object.   These two functions have a major restriction in that
+    they are **mutually exclusive** based on the type of object to be returned.
+    The ``JSON_QUERY`` function **only** returns a JSON dictionary or list,
+    but not an individual string, numeric, or boolean element; the
+    ``JSON_VALUE`` function **only** returns an individual string, numeric,
+    or boolean element.   **both functions either return NULL or raise
+    an error if they are not used against the correct expected value**.
+
+    To handle this awkward requirement, indexed access rules are as follows:
+
+    1. When extracting a sub element from a JSON that is itself a JSON
+       dictionary or list, the :meth:`_types.JSON.Comparator.as_json` accessor
+       should be used::
+
+            stmt = select(
+                data_table.c.data["some key"].as_json()
+            ).where(
+                data_table.c.data["some key"].as_json() == {"sub": "structure"}
+            )
+
+    2. When extracting a sub element from a JSON that is a plain boolean,
+       string, integer, or float, use the appropriate method among
+       :meth:`_types.JSON.Comparator.as_boolean`,
+       :meth:`_types.JSON.Comparator.as_string`,
+       :meth:`_types.JSON.Comparator.as_integer`,
+       :meth:`_types.JSON.Comparator.as_float`::
+
+            stmt = select(
+                data_table.c.data["some key"].as_string()
+            ).where(
+                data_table.c.data["some key"].as_string() == "some string"
+            )
+
+    .. versionadded:: 1.4
+
+
+    """
+
+    # note there was a result processor here that was looking for "number",
+    # but none of the tests seem to exercise it.
+
+
+# Note: these objects currently match exactly those of MySQL, however since
+# these are not generalizable to all JSON implementations, remain separately
+# implemented for each dialect.
+class _FormatTypeMixin(object):
+    def _format_value(self, value):
+        raise NotImplementedError()
+
+    def bind_processor(self, dialect):
+        super_proc = self.string_bind_processor(dialect)
+
+        def process(value):
+            value = self._format_value(value)
+            if super_proc:
+                value = super_proc(value)
+            return value
+
+        return process
+
+    def literal_processor(self, dialect):
+        super_proc = self.string_literal_processor(dialect)
+
+        def process(value):
+            value = self._format_value(value)
+            if super_proc:
+                value = super_proc(value)
+            return value
+
+        return process
+
+
+class JSONIndexType(_FormatTypeMixin, sqltypes.JSON.JSONIndexType):
+    def _format_value(self, value):
+        if isinstance(value, int):
+            value = "$[%s]" % value
+        else:
+            value = '$."%s"' % value
+        return value
+
+
+class JSONPathType(_FormatTypeMixin, sqltypes.JSON.JSONPathType):
+    def _format_value(self, value):
+        return "$%s" % (
+            "".join(
+                [
+                    "[%s]" % elem if isinstance(elem, int) else '."%s"' % elem
+                    for elem in value
+                ]
+            )
+        )
diff --git a/lib/sqlalchemy/dialects/mysql/base.py b/lib/sqlalchemy/dialects/mysql/base.py
index 34afc81a7..1d032b600 100644
--- a/lib/sqlalchemy/dialects/mysql/base.py
+++ b/lib/sqlalchemy/dialects/mysql/base.py
@@ -1417,15 +1417,20 @@ class MySQLCompiler(compiler.SQLCompiler):
             # explicitly return true/false constants
             type_expression = "WHEN true THEN true ELSE false"
         elif binary.type._type_affinity is sqltypes.String:
-            # this fails with a JSON value that's a four byte unicode
+            # (gord): this fails with a JSON value that's a four byte unicode
             # string.  SQLite has the same problem at the moment
+            # (zzzeek): I'm not really sure.  let's take a look at a test case
+            # that hits each backend and maybe make a requires rule for it?
             type_expression = "ELSE JSON_UNQUOTE(JSON_EXTRACT(%s, %s))" % (
                 self.process(binary.left, **kw),
                 self.process(binary.right, **kw),
             )
         else:
             # other affinity....this is not expected right now
-            type_expression = "ELSE JSON_EXTRACT(%s, %s)"
+            type_expression = "ELSE JSON_EXTRACT(%s, %s)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
 
         return case_expression + " " + type_expression + " END"
 
diff --git a/lib/sqlalchemy/dialects/mysql/json.py b/lib/sqlalchemy/dialects/mysql/json.py
index 640e19929..a1c8258b0 100644
--- a/lib/sqlalchemy/dialects/mysql/json.py
+++ b/lib/sqlalchemy/dialects/mysql/json.py
@@ -16,6 +16,14 @@ class JSON(sqltypes.JSON):
     MySQL supports JSON as of version 5.7.
     MariaDB supports JSON (as an alias for LONGTEXT) as of version 10.2.
 
+    :class:`_mysql.JSON` is used automatically whenever the base
+    :class:`_types.JSON` datatype is used against a MySQL or MariaDB backend.
+
+    .. seealso::
+
+        :class:`_types.JSON` - main documenation for the generic
+        cross-platform JSON datatype.
+
     The :class:`.mysql.JSON` type supports persistence of JSON values
     as well as the core index operations provided by :class:`_types.JSON`
     datatype, by adapting the operations to render the ``JSON_EXTRACT``
diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index 255f1af21..fbf61dd5f 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -97,12 +97,18 @@ class JSONPathType(sqltypes.JSON.JSONPathType):
 class JSON(sqltypes.JSON):
     """Represent the PostgreSQL JSON type.
 
-    This type is a specialization of the Core-level :class:`_types.JSON`
-    type.   Be sure to read the documentation for :class:`_types.JSON` for
-    important tips regarding treatment of NULL values and ORM use.
+    :class:`_postgresql.JSON` is used automatically whenever the base
+    :class:`_types.JSON` datatype is used against a PostgreSQL backend,
+    however base :class:`_types.JSON` datatype does not provide Python
+    accessors for PostgreSQL-specific comparison methods such as
+    :meth:`_postgresql.JSON.Comparator.astext`; additionally, to use
+    PostgreSQL ``JSONB``, the :class:`_postgresql.JSONB` datatype should
+    be used explicitly.
 
-    .. versionchanged:: 1.1 :class:`_postgresql.JSON` is now a PostgreSQL-
-       specific specialization of the new :class:`_types.JSON` type.
+    .. seealso::
+
+        :class:`_types.JSON` - main documenation for the generic
+        cross-platform JSON datatype.
 
     The operators provided by the PostgreSQL version of :class:`_types.JSON`
     include:
@@ -167,6 +173,9 @@ class JSON(sqltypes.JSON):
 
         :class:`_postgresql.JSONB`
 
+    .. versionchanged:: 1.1 :class:`_postgresql.JSON` is now a PostgreSQL-
+       specific specialization of the new :class:`_types.JSON` type.
+
     """  # noqa
 
     astext_type = sqltypes.Text()
diff --git a/lib/sqlalchemy/dialects/sqlite/json.py b/lib/sqlalchemy/dialects/sqlite/json.py
index 775f557f8..fadec3ce3 100644
--- a/lib/sqlalchemy/dialects/sqlite/json.py
+++ b/lib/sqlalchemy/dialects/sqlite/json.py
@@ -9,6 +9,14 @@ class JSON(sqltypes.JSON):
     `loadable extension <https://www.sqlite.org/loadext.html>`_ and as such
     may not be available, or may require run-time loading.
 
+    :class:`_sqlite.JSON` is used automatically whenever the base
+    :class:`_types.JSON` datatype is used against a SQLite backend.
+
+    .. seealso::
+
+        :class:`_types.JSON` - main documenation for the generic
+        cross-platform JSON datatype.
+
     The :class:`_sqlite.JSON` type supports persistence of JSON values
     as well as the core index operations provided by :class:`_types.JSON`
     datatype, by adapting the operations to render the ``JSON_EXTRACT``
@@ -16,11 +24,9 @@ class JSON(sqltypes.JSON):
     Extracted values are quoted in order to ensure that the results are
     always JSON string values.
 
-    .. versionadded:: 1.3
 
-    .. seealso::
+    .. versionadded:: 1.3
 
-        JSON1_
 
     .. _JSON1: https://www.sqlite.org/json1.html
 
diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
index fd85d6d30..f1063b71a 100644
--- a/lib/sqlalchemy/sql/sqltypes.py
+++ b/lib/sqlalchemy/sql/sqltypes.py
@@ -2034,11 +2034,17 @@ class JSON(Indexable, TypeEngine):
        JSON types.  Since it supports JSON SQL operations, it only
        works on backends that have an actual JSON type, currently:
 
-       * PostgreSQL
+       * PostgreSQL - see :class:`_postgresql.JSON` and
+         :class:`_postgresql.JSONB` for backend-specific notes
 
-       * MySQL as of version 5.7 (MariaDB as of the 10.2 series does not)
+       * MySQL as of version 5.7 (MariaDB as of the 10.2 series does not) - see
+         :class:`_mysql.JSON` for backend-specific notes
 
-       * SQLite as of version 3.9
+       * SQLite as of version 3.9 - see :class:`_sqlite.JSON` for
+         backend-specific notes
+
+       * Microsoft SQL Server 2016 and later - see :class:`_mssql.JSON` for
+         backend-specific notes
 
     :class:`_types.JSON` is part of the Core in support of the growing
     popularity of native JSON datatypes.
@@ -2452,7 +2458,14 @@ class JSON(Indexable, TypeEngine):
         def as_json(self):
             """Cast an indexed value as JSON.
 
-            This is the default behavior of indexed elements in any case.
+            e.g.::
+
+                stmt = select([
+                    mytable.c.json_column['some_data'].as_json()
+                ])
+
+            This is typically the default behavior of indexed elements in any
+            case.
 
             Note that comparison of full JSON structures may not be
             supported by all backends.
diff --git a/lib/sqlalchemy/testing/requirements.py b/lib/sqlalchemy/testing/requirements.py
index 3d3980b30..301c9ef84 100644
--- a/lib/sqlalchemy/testing/requirements.py
+++ b/lib/sqlalchemy/testing/requirements.py
@@ -804,6 +804,16 @@ class SuiteRequirements(Requirements):
         return exclusions.open()
 
     @property
+    def legacy_unconditional_json_extract(self):
+        """Backend has a JSON_EXTRACT or similar function that returns a
+        valid JSON string in all cases.
+
+        Used to test a legacy feature and is not needed.
+
+        """
+        return exclusions.closed()
+
+    @property
     def precision_numerics_general(self):
         """target backend has general support for moderately high-precision
         numerics."""
diff --git a/lib/sqlalchemy/testing/suite/test_types.py b/lib/sqlalchemy/testing/suite/test_types.py
index 6a390231b..9a2fdf95a 100644
--- a/lib/sqlalchemy/testing/suite/test_types.py
+++ b/lib/sqlalchemy/testing/suite/test_types.py
@@ -774,8 +774,21 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
 
         eq_(row, (data_element,))
 
-    def _index_fixtures(fn):
-        fn = testing.combinations(
+    def _index_fixtures(include_comparison):
+
+        if include_comparison:
+            # basically SQL Server and MariaDB can kind of do json
+            # comparison, MySQL, PG and SQLite can't.  not worth it.
+            json_elements = []
+        else:
+            json_elements = [
+                ("json", {"foo": "bar"}),
+                ("json", ["one", "two", "three"]),
+                (None, {"foo": "bar"}),
+                (None, ["one", "two", "three"]),
+            ]
+
+        elements = [
             ("boolean", True),
             ("boolean", False),
             ("boolean", None),
@@ -793,14 +806,16 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
             ("integer", None),
             ("float", 28.5),
             ("float", None),
-            # TODO: how to test for comaprison
-            #        ("json", {"foo": "bar"}),
-            id_="sa",
-        )(fn)
+        ] + json_elements
+
+        def decorate(fn):
+            fn = testing.combinations(id_="sa", *elements)(fn)
 
-        return fn
+            return fn
 
-    @_index_fixtures
+        return decorate
+
+    @_index_fixtures(False)
     def test_index_typed_access(self, datatype, value):
         data_table = self.tables.data_table
         data_element = {"key1": value}
@@ -815,14 +830,16 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
             )
 
             expr = data_table.c.data["key1"]
-            expr = getattr(expr, "as_%s" % datatype)()
+
+            if datatype:
+                expr = getattr(expr, "as_%s" % datatype)()
 
             roundtrip = conn.scalar(select(expr))
             eq_(roundtrip, value)
             if util.py3k:  # skip py2k to avoid comparing unicode to str etc.
                 is_(type(roundtrip), type(value))
 
-    @_index_fixtures
+    @_index_fixtures(True)
     def test_index_typed_comparison(self, datatype, value):
         data_table = self.tables.data_table
         data_element = {"key1": value}
@@ -837,14 +854,15 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
             )
 
             expr = data_table.c.data["key1"]
-            expr = getattr(expr, "as_%s" % datatype)()
+            if datatype:
+                expr = getattr(expr, "as_%s" % datatype)()
 
             row = conn.execute(select(expr).where(expr == value)).first()
 
             # make sure we get a row even if value is None
             eq_(row, (value,))
 
-    @_index_fixtures
+    @_index_fixtures(True)
     def test_path_typed_comparison(self, datatype, value):
         data_table = self.tables.data_table
         data_element = {"key1": {"subkey1": value}}
@@ -859,7 +877,9 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
             )
 
             expr = data_table.c.data[("key1", "subkey1")]
-            expr = getattr(expr, "as_%s" % datatype)()
+
+            if datatype:
+                expr = getattr(expr, "as_%s" % datatype)()
 
             row = conn.execute(select(expr).where(expr == value)).first()
 
@@ -1033,14 +1053,17 @@ class JSONTest(_LiteralRoundTripFixture, fixtures.TablesTest):
         )
 
 
-class JSONStringCastIndexTest(_LiteralRoundTripFixture, fixtures.TablesTest):
+class JSONLegacyStringCastIndexTest(
+    _LiteralRoundTripFixture, fixtures.TablesTest
+):
     """test JSON index access with "cast to string", which we have documented
     for a long time as how to compare JSON values, but is ultimately not
-    reliable in all cases.
+    reliable in all cases.   The "as_XYZ()" comparators should be used
+    instead.
 
     """
 
-    __requires__ = ("json_type",)
+    __requires__ = ("json_type", "legacy_unconditional_json_extract")
     __backend__ = True
 
     datatype = JSON
@@ -1135,13 +1158,13 @@ class JSONStringCastIndexTest(_LiteralRoundTripFixture, fixtures.TablesTest):
         # "cannot extract array element from a non-array", which is
         # fixed in 9.4 but may exist in 9.3
         self._test_index_criteria(
-            and_(name == "r4", cast(col[1], String) == '"two"'), "r4"
+            and_(name == "r4", cast(col[1], String) == '"two"',), "r4",
         )
 
     def test_string_cast_crit_mixed_path(self):
         col = self.tables.data_table.c["data"]
         self._test_index_criteria(
-            cast(col[("key3", 1, "six")], String) == '"seven"', "r3"
+            cast(col[("key3", 1, "six")], String) == '"seven"', "r3",
         )
 
     def test_string_cast_crit_string_path(self):
@@ -1157,7 +1180,8 @@ class JSONStringCastIndexTest(_LiteralRoundTripFixture, fixtures.TablesTest):
         col = self.tables.data_table.c["data"]
 
         self._test_index_criteria(
-            and_(name == "r6", cast(col["b"], String) == '"some value"'), "r6"
+            and_(name == "r6", cast(col["b"], String) == '"some value"',),
+            "r6",
         )
 
 
@@ -1165,7 +1189,7 @@ __all__ = (
     "UnicodeVarcharTest",
     "UnicodeTextTest",
     "JSONTest",
-    "JSONStringCastIndexTest",
+    "JSONLegacyStringCastIndexTest",
     "DateTest",
     "DateTimeTest",
     "TextTest",