fix(django-google-spanner): support JSONField read and write operations

packages/django-google-spanner/django_spanner/__init__.py

@@ -13,7 +13,6 @@
 
 import django
 from django.db import DEFAULT_DB_ALIAS
-from django.db.models import JSONField
 from django.db.models.fields import (
     NOT_PROVIDED,
     AutoField,
@@ -27,7 +26,6 @@
 from google.api_core.datetime_helpers import (
     DatetimeWithNanoseconds,
 )
-from google.cloud.spanner_v1 import JsonObject
 
 from .functions import register_functions
 from .lookups import register_lookups
@@ -84,17 +82,6 @@ def autofield_init(self, *args, **kwargs):
 SmallAutoField.__init__ = autofield_init
 BigAutoField.__init__ = autofield_init
 
-
-def get_prep_value(self, value):
-    # Json encoding and decoding for spanner is done in python-spanner.
-    if not isinstance(value, JsonObject) and isinstance(value, dict):
-        return JsonObject(value)
-
-    return value
-
-
-JSONField.get_prep_value = get_prep_value
-
 old_datetimewithnanoseconds_eq = getattr(DatetimeWithNanoseconds, "__eq__", None)
 
 

packages/django-google-spanner/django_spanner/compiler.py

@@ -5,6 +5,7 @@
 # https://developers.google.com/open-source/licenses/bsd
 
 from django.core.exceptions import EmptyResultSet
+from django.db.models.fields.json import JSONField
 from django.db.models.sql.compiler import (
     SQLAggregateCompiler as BaseSQLAggregateCompiler,
 )
@@ -31,6 +32,28 @@ class SQLCompiler(BaseSQLCompiler):
     functionality.
     """
 
+    def get_converters(self, expressions):
+        converters = {}
+        for i, expression in enumerate(expressions):
+            if expression:
+                backend_converters = self.connection.ops.get_db_converters(expression)
+                field_converters = expression.get_db_converters(self.connection)
+                output_field = getattr(expression, "output_field", None)
+                if (
+                    self.connection.ops.convert_jsonfield_value in backend_converters
+                    and isinstance(output_field, JSONField)
+                    and type(output_field).from_db_value is JSONField.from_db_value
+                ):
+                    field_converters = [
+                        converter
+                        for converter in field_converters
+                        if converter != output_field.from_db_value
+                    ]
+                convs = backend_converters + field_converters
+                if convs:
+                    converters[i] = (convs, expression)
+        return converters
+
     def get_combinator_sql(self, combinator, all):
         """Override the native Django method.
 

packages/django-google-spanner/django_spanner/features.py

@@ -57,16 +57,16 @@ def introspected_field_types(self):
     supports_timezones = False
     supports_transactions = True
     if USE_EMULATOR:
-        # Emulator does not support json.
-        supports_json_field = False
         # Emulator does not support check constraints.
         supports_column_check_constraints = False
         supports_table_check_constraints = False
     else:
         supports_column_check_constraints = True
         supports_table_check_constraints = True
-        supports_json_field = True
-    supports_primitives_in_json_field = False
+    has_native_json_field = True
+    supports_json_field = True
+    supports_primitives_in_json_field = True
+    supports_json_field_contains = False
     supports_composite_primary_keys = True
     # Spanner does not support order by null modifiers.
     supports_order_by_nulls_modifier = False
@@ -107,6 +107,7 @@ def introspected_field_types(self):
         "constraints.tests.CheckConstraintTests.test_validate_boolean_expressions",
         "constraints.tests.CheckConstraintTests.test_validate_custom_error",
         "constraints.tests.CheckConstraintTests.test_validate_generated_field_stored",
+        "constraints.tests.CheckConstraintTests.test_validate_jsonfield_exact",
         "constraints.tests.CheckConstraintTests.test_validate_pk_field",
         "constraints.tests.UniqueConstraintTests.test_validate_case_when",
         "constraints.tests.UniqueConstraintTests.test_validate_expression_condition",
@@ -132,6 +133,17 @@ def introspected_field_types(self):
         "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_lookup_name_sql_injection",
         "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_time_comparison",
         "db_functions.datetime.test_extract_trunc.DateFunctionWithTimeZoneTests.test_trunc_timezone_applied_before_truncation",
+        "db_functions.json.test_json_array.JSONArrayObjectTests.test_nested_json_array_object",
+        "db_functions.json.test_json_array.JSONArrayObjectTests.test_nested_json_object_array",
+        "db_functions.json.test_json_array.JSONArrayObjectTests.test_order_by_nested_key",
+        "db_functions.json.test_json_array.JSONArrayTests.test_basic",
+        "db_functions.json.test_json_array.JSONArrayTests.test_empty",
+        "db_functions.json.test_json_array.JSONArrayTests.test_expressions",
+        "db_functions.json.test_json_array.JSONArrayTests.test_nested_empty_json_array",
+        "db_functions.json.test_json_array.JSONArrayTests.test_nested_json_array",
+        "db_functions.json.test_json_array.JSONArrayTests.test_order_by_key",
+        "db_functions.json.test_json_array.JSONArrayTests.test_order_by_nested_key",
+        "db_functions.json.test_json_array.JSONArrayTests.test_textfield",
         "db_functions.json.test_json_object.JSONObjectTests.test_basic",
         "db_functions.json.test_json_object.JSONObjectTests.test_empty",
         "db_functions.json.test_json_object.JSONObjectTests.test_expressions",
@@ -1512,8 +1524,19 @@ def introspected_field_types(self):
             "model_fields.test_jsonfield.TestSerialization.test_xml_serialization",
             "model_fields.test_jsonfield.TestSerialization.test_xml_serialization",
             "model_fields.test_jsonfield.TestSerialization.test_xml_serialization",
+            "model_fields.test_jsonfield.TestQuerying.test_deep_values",
+            "model_fields.test_jsonfield.TestQuerying.test_expression_wrapper_key_transform",
+            "model_fields.test_jsonfield.TestQuerying.test_has_key_literal_lookup",
+            "model_fields.test_jsonfield.TestQuerying.test_has_key_number",
+            "model_fields.test_jsonfield.TestQuerying.test_icontains",
+            "model_fields.test_jsonfield.TestQuerying.test_key_sql_injection_escape",
+            "model_fields.test_jsonfield.TestQuerying.test_key_transform_raw_expression",
+            "model_fields.test_jsonfield.TestQuerying.test_lookups_special_chars",
+            "model_fields.test_jsonfield.TestQuerying.test_lookups_special_chars_double_quotes",
+            "model_fields.test_jsonfield.TestQuerying.test_nested_key_transform_raw_expression",
             "model_fields.test_uuid.TestSerialization.test_dumping",
             "queries.test_bulk_update.BulkUpdateNoteTests.test_unsaved_models",
+            "queries.test_bulk_update.BulkUpdateTests.test_json_field_sql_null",
             "queries.test_bulk_update.BulkUpdateTests.test_large_batch",
             "queries.test_bulk_update.BulkUpdateTests.test_unsaved_parent",
             "queries.test_qs_combinators.QuerySetSetOperationTests.test_difference_with_values",
@@ -2283,6 +2306,7 @@ def introspected_field_types(self):
 
     django_6_0_skip_tests = (
         # Spanner uses random int64 IDs; test assumes monotonic ordering matching insertion order.
+        "prefetch_related.tests.PrefetchRelatedMTICacheTests.test_grandparent_m2m_available_in_child",
         "prefetch_related.tests.PrefetchRelatedMTICacheTests.test_parent_m2m_available_in_child",
         # Client-side AutoField ID generation sets pk before save; _is_pk_set() triggers refresh_from_db instead of AttributeError.
         "defer_regress.tests.DeferCopyInstanceTests.test_bulk_create",
@@ -2303,6 +2327,10 @@ def introspected_field_types(self):
         "bulk_create.tests.BulkCreateTransactionTests.test_objs_with_and_without_pk",
         # Tie-breaker ordering on pk assumes sequential integer IDs; Spanner uses random IDs.
         "ordering.tests.OrderingTests.test_order_by_case_when_constant_value",
+        # Spanner GoogleSQL JSONPath does not support negative array indices (e.g. $[-1]).
+        "model_fields.test_jsonfield.TestQuerying.test_deep_negative_lookup_array",
+        "model_fields.test_jsonfield.TestQuerying.test_deep_negative_lookup_mixed",
+        "model_fields.test_jsonfield.TestQuerying.test_shallow_list_negative_lookup",
     )
 
     if django.VERSION >= (6, 0):

packages/django-google-spanner/django_spanner/lookups.py

@@ -4,6 +4,11 @@
 # license that can be found in the LICENSE file or at
 # https://developers.google.com/open-source/licenses/bsd
 
+from django.db.models.fields.json import (
+    HasKeyLookup,
+    KeyTextTransform,
+    KeyTransform,
+)
 from django.db.models.lookups import (
     Contains,
     EndsWith,
@@ -263,6 +268,42 @@ def cast_param_to_float(self, compiler, connection):
     return sql, tuple(params)
 
 
+def _compile_json_path(connection, key_transforms):
+    if hasattr(connection.ops, "compile_json_path"):
+        return connection.ops.compile_json_path(key_transforms)
+    from django.db.models.fields.json import compile_json_path
+
+    return compile_json_path(key_transforms)
+
+
+def _quote_json_path(json_path):
+    """Inline a compiled JSONPath as a safely escaped GoogleSQL string literal."""
+    escaped = json_path.replace("\\", "\\\\").replace("'", "\\'").replace("%", "%%")
+    return f"'{escaped}'"
+
+
+def key_transform_as_spanner(self, compiler, connection):
+    lhs, params, key_transforms = self.preprocess_lhs(compiler, connection)
+    json_path = _quote_json_path(_compile_json_path(connection, key_transforms))
+    return f"JSON_QUERY({lhs}, {json_path})", tuple(params)
+
+
+def key_text_transform_as_spanner(self, compiler, connection):
+    lhs, params, key_transforms = self.preprocess_lhs(compiler, connection)
+    json_path = _quote_json_path(_compile_json_path(connection, key_transforms))
+    return f"JSON_VALUE({lhs}, {json_path})", tuple(params)
+
+
+def has_key_lookup_as_spanner(self, compiler, connection):
+    sql_parts = []
+    params = []
+    for lhs_sql, lhs_params, rhs_json_path in self._as_sql_parts(compiler, connection):
+        json_path = _quote_json_path(rhs_json_path)
+        sql_parts.append(f"JSON_QUERY({lhs_sql}, {json_path}) IS NOT NULL")
+        params.extend(lhs_params)
+    return self._combine_sql_parts(sql_parts), tuple(params)
+
+
 def register_lookups():
     """Registers the above methods with the corersponding Django classes."""
     Contains.as_spanner = contains
@@ -279,3 +320,6 @@ def register_lookups():
     GreaterThanOrEqual.as_spanner = cast_param_to_float
     LessThan.as_spanner = cast_param_to_float
     LessThanOrEqual.as_spanner = cast_param_to_float
+    KeyTransform.as_spanner = key_transform_as_spanner
+    KeyTextTransform.as_spanner = key_text_transform_as_spanner
+    HasKeyLookup.as_spanner = has_key_lookup_as_spanner

packages/django-google-spanner/django_spanner/operations.py

@@ -4,6 +4,7 @@
 # license that can be found in the LICENSE file or at
 # https://developers.google.com/open-source/licenses/bsd
 
+import json
 import os
 import re
 from base64 import b64decode
@@ -12,14 +13,43 @@
 
 from django.conf import settings
 from django.db.backends.base.operations import BaseDatabaseOperations
+from django.db.models.fields.json import JSONField
 from django.db.utils import DatabaseError
 from django.utils import timezone
 from django.utils.duration import duration_microseconds
 from google.cloud.spanner_dbapi.parse_utils import (
+    TYPES_MAP,
     DateStr,
     TimestampStr,
     escape_name,
 )
+from google.cloud.spanner_v1 import JsonObject, param_types
+
+
+class _JsonString(JsonObject):
+    """Wrap a pre-serialized JSON string for spanner_dbapi without re-parsing."""
+
+    def __init__(self, raw_json):
+        self._is_null = False
+        self._is_array = False
+        self._is_scalar_value = False
+        self._raw_json = raw_json
+
+    def serialize(self):
+        return self._raw_json
+
+    def __repr__(self):
+        return f"_JsonString({self._raw_json!r})"
+
+    def __eq__(self, other):
+        if isinstance(other, _JsonString):
+            return self._raw_json == other._raw_json
+        if isinstance(other, JsonObject):
+            return self._raw_json == other.serialize()
+        return False
+
+
+TYPES_MAP[_JsonString] = param_types.JSON
 
 
 class DatabaseOperations(BaseDatabaseOperations):
@@ -291,6 +321,13 @@ def adapt_timefield_value(self, value):
             "0001-01-01T" + value.isoformat(timespec="microseconds") + "Z"
         )
 
+    def adapt_json_value(self, value, encoder):
+        """Wrap Django's serialized JSON string for Spanner JSON parameter binding."""
+        raw_json = super().adapt_json_value(value, encoder)
+        if raw_json is None:
+            return None
+        return _JsonString(raw_json)
+
     def get_db_converters(self, expression):
         """Get a list of functions needed to convert field data.
 
@@ -310,6 +347,8 @@ def get_db_converters(self, expression):
             converters.append(self.convert_binaryfield_value)
         elif internal_type == "UUIDField":
             converters.append(self.convert_uuidfield_value)
+        elif internal_type == "JSONField":
+            converters.append(self.convert_jsonfield_value)
         return converters
 
     def convert_binaryfield_value(self, value, expression, connection):
@@ -406,6 +445,35 @@ def convert_uuidfield_value(self, value, expression, connection):
             value = UUID(value)
         return value
 
+    def convert_jsonfield_value(self, value, expression, connection):
+        """Convert Spanner JsonObject received from DB to a Python value for Django."""
+        if value is None:
+            return None
+        output_field = getattr(expression, "output_field", None)
+        has_custom_from_db_value = (
+            isinstance(output_field, JSONField)
+            and type(output_field).from_db_value is not JSONField.from_db_value
+        )
+        decoder = getattr(output_field, "decoder", None)
+        if isinstance(value, JsonObject):
+            if value._is_null:
+                return None
+            if has_custom_from_db_value:
+                return value.serialize()
+            if decoder is not None:
+                return json.loads(value.serialize(), cls=decoder)
+            if value._is_array:
+                return value._array_value
+            if value._is_scalar_value:
+                return value._simple_value
+            return dict(value)
+        if isinstance(value, str) and not has_custom_from_db_value:
+            try:
+                return json.loads(value, cls=decoder)
+            except json.JSONDecodeError:
+                return value
+        return value
+
     def date_extract_sql(self, lookup_type, field_name, params=None):
         """Extract date from the lookup.
 

packages/django-google-spanner/tests/mockserver_tests/models.py

@@ -12,9 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from django.core.serializers.json import DjangoJSONEncoder
 from django.db import models
 
 
 class Singer(models.Model):
     first_name = models.CharField(max_length=200)
     last_name = models.CharField(max_length=200)
+
+
+class Detail(models.Model):
+    value = models.JSONField(null=True, encoder=DjangoJSONEncoder)