Add type accessors for JSON indexed/pathed element access

Added new accessors to expressions of type :class:`.JSON` to allow for
specific datatype access and comparison, covering strings, integers,
numeric, boolean elements.   This revises the documented approach of
CASTing to string when comparing values, instead adding specific
functionality into the PostgreSQL, SQlite, MySQL dialects to reliably
deliver these basic types in all cases.

The change also delivers a new feature to the test exclusions
system so that combinations and exclusions can be used together.

Fixes: #4276
Change-Id: Ica5a926c060feb40a0a7cd60b9d6e061d7825728

4 files changed, 97 insertions, 13 deletions

diff --git a/lib/sqlalchemy/dialects/mysql/base.py b/lib/sqlalchemy/dialects/mysql/base.py
index 05edb6310..fb123bc0f 100644
--- a/lib/sqlalchemy/dialects/mysql/base.py
+++ b/lib/sqlalchemy/dialects/mysql/base.py
@@ -1207,17 +1207,63 @@ class MySQLCompiler(compiler.SQLCompiler):
     def visit_sysdate_func(self, fn, **kw):
         return "SYSDATE()"
 
-    def visit_json_getitem_op_binary(self, binary, operator, **kw):
-        return "JSON_EXTRACT(%s, %s)" % (
+    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_EXTRACT(%s, %s)" % (
+                self.process(binary.left, **kw),
+                self.process(binary.right, **kw),
+            )
+
+        # for non-JSON, MySQL doesn't handle JSON null at all so it has to
+        # be explicit
+        case_expression = "CASE JSON_EXTRACT(%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_EXTRACT(%s, %s) AS SIGNED INTEGER)"
+                % (
+                    self.process(binary.left, **kw),
+                    self.process(binary.right, **kw),
+                )
+            )
+        elif binary.type._type_affinity is sqltypes.Numeric:
+            # FLOAT / REAL not added in MySQL til 8.0.17
+            type_expression = (
+                "ELSE CAST(JSON_EXTRACT(%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 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
+            # string.  SQLite has the same problem at the moment
+            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)"
+
+        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 "JSON_EXTRACT(%s, %s)" % (
-            self.process(binary.left, **kw),
-            self.process(binary.right, **kw),
-        )
+        return self._render_json_extract_from_binary(binary, operator, **kw)
 
     def visit_on_duplicate_key_update(self, on_duplicate, **kw):
         if on_duplicate._parameter_ordering:
diff --git a/lib/sqlalchemy/dialects/postgresql/base.py b/lib/sqlalchemy/dialects/postgresql/base.py
index d6fd2623b..6d97033d0 100644
--- a/lib/sqlalchemy/dialects/postgresql/base.py
+++ b/lib/sqlalchemy/dialects/postgresql/base.py
@@ -1559,13 +1559,36 @@ class PGCompiler(compiler.SQLCompiler):
             self.process(element.stop, **kw),
         )
 
-    def visit_json_getitem_op_binary(self, binary, operator, **kw):
+    def visit_json_getitem_op_binary(
+        self, binary, operator, _cast_applied=False, **kw
+    ):
+        if (
+            not _cast_applied
+            and binary.type._type_affinity is not sqltypes.JSON
+        ):
+            kw["_cast_applied"] = True
+            return self.process(sql.cast(binary, binary.type), **kw)
+
         kw["eager_grouping"] = True
-        return self._generate_generic_binary(binary, " -> ", **kw)
 
-    def visit_json_path_getitem_op_binary(self, binary, operator, **kw):
+        return self._generate_generic_binary(
+            binary, " -> " if not _cast_applied else " ->> ", **kw
+        )
+
+    def visit_json_path_getitem_op_binary(
+        self, binary, operator, _cast_applied=False, **kw
+    ):
+        if (
+            not _cast_applied
+            and binary.type._type_affinity is not sqltypes.JSON
+        ):
+            kw["_cast_applied"] = True
+            return self.process(sql.cast(binary, binary.type), **kw)
+
         kw["eager_grouping"] = True
-        return self._generate_generic_binary(binary, " #> ", **kw)
+        return self._generate_generic_binary(
+            binary, " #> " if not _cast_applied else " #>> ", **kw
+        )
 
     def visit_getitem_binary(self, binary, operator, **kw):
         return "%s[%s]" % (
diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index b8e3ac23d..ef48f3d86 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -123,11 +123,17 @@ class JSON(sqltypes.JSON):
 
         data_table.c.data['some key'].astext == 'some value'
 
+      Note that equivalent functionality is available via the
+      :attr:`.JSON.Comparator.as_string` accessor.
+
     * Index operations with CAST
       (equivalent to ``CAST(col ->> ['some key'] AS <type>)``)::
 
         data_table.c.data['some key'].astext.cast(Integer) == 5
 
+      Note that equivalent functionality is available via the
+      :attr:`.JSON.Comparator.as_integer` and similar accessors.
+
     * Path index operations (the ``#>`` operator)::
 
         data_table.c.data[('key_1', 'key_2', 5, ..., 'key_n')]
@@ -214,7 +220,6 @@ class JSON(sqltypes.JSON):
                 :meth:`.ColumnElement.cast`
 
             """
-
             if isinstance(self.expr.right.type, sqltypes.JSON.JSONPathType):
                 return self.expr.left.operate(
                     JSONPATH_ASTEXT,
diff --git a/lib/sqlalchemy/dialects/sqlite/base.py b/lib/sqlalchemy/dialects/sqlite/base.py
index 02d44a260..2685a9243 100644
--- a/lib/sqlalchemy/dialects/sqlite/base.py
+++ b/lib/sqlalchemy/dialects/sqlite/base.py
@@ -1012,13 +1012,23 @@ class SQLiteCompiler(compiler.SQLCompiler):
         )
 
     def visit_json_getitem_op_binary(self, binary, operator, **kw):
-        return "JSON_QUOTE(JSON_EXTRACT(%s, %s))" % (
+        if binary.type._type_affinity is sqltypes.JSON:
+            expr = "JSON_QUOTE(JSON_EXTRACT(%s, %s))"
+        else:
+            expr = "JSON_EXTRACT(%s, %s)"
+
+        return expr % (
             self.process(binary.left, **kw),
             self.process(binary.right, **kw),
         )
 
     def visit_json_path_getitem_op_binary(self, binary, operator, **kw):
-        return "JSON_QUOTE(JSON_EXTRACT(%s, %s))" % (
+        if binary.type._type_affinity is sqltypes.JSON:
+            expr = "JSON_QUOTE(JSON_EXTRACT(%s, %s))"
+        else:
+            expr = "JSON_EXTRACT(%s, %s)"
+
+        return expr % (
             self.process(binary.left, **kw),
             self.process(binary.right, **kw),
         )