diff --git a/packages/django-google-spanner/django_spanner/__init__.py b/packages/django-google-spanner/django_spanner/__init__.py index fc8cb56b2104..49d9f996d34c 100644 --- a/packages/django-google-spanner/django_spanner/__init__.py +++ b/packages/django-google-spanner/django_spanner/__init__.py @@ -93,7 +93,27 @@ def get_prep_value(self, value): return value +old_jsonfield_from_db_value = getattr(JSONField, "from_db_value", None) + + +def jsonfield_from_db_value(self, value, expression, connection): + if isinstance(value, JsonObject): + if value._is_null: + return None + if value._is_array: + return value._array_value + if value._is_scalar_value: + return value._simple_value + return dict(value) + if isinstance(value, (dict, list)): + return value + if old_jsonfield_from_db_value: + return old_jsonfield_from_db_value(self, value, expression, connection) + return value + + JSONField.get_prep_value = get_prep_value +JSONField.from_db_value = jsonfield_from_db_value old_datetimewithnanoseconds_eq = getattr(DatetimeWithNanoseconds, "__eq__", None) diff --git a/packages/django-google-spanner/django_spanner/operations.py b/packages/django-google-spanner/django_spanner/operations.py index 5a5b8e3e4612..4ed0f14c4e0e 100644 --- a/packages/django-google-spanner/django_spanner/operations.py +++ b/packages/django-google-spanner/django_spanner/operations.py @@ -20,6 +20,7 @@ TimestampStr, escape_name, ) +from google.cloud.spanner_v1 import JsonObject class DatabaseOperations(BaseDatabaseOperations): @@ -770,3 +771,12 @@ def savepoint_rollback_sql(self, sid): Return the SQL for rolling back to a savepoint. """ return "SELECT 1" + + def adapt_json_value(self, value, encoder): + if value is None: + return None + if isinstance(value, JsonObject): + if value._is_null or value._is_array or value._is_scalar_value: + return value + value = dict(value) + return JsonObject.from_str(super().adapt_json_value(value, encoder)) diff --git a/packages/django-google-spanner/tests/unit/django_spanner/test_operations.py b/packages/django-google-spanner/tests/unit/django_spanner/test_operations.py index 520d268bce10..7d9aa7bc6fa1 100644 --- a/packages/django-google-spanner/tests/unit/django_spanner/test_operations.py +++ b/packages/django-google-spanner/tests/unit/django_spanner/test_operations.py @@ -347,3 +347,37 @@ def test_datetime_cast_sql(self): "TIMESTAMP(FORMAT_TIMESTAMP('%Y-%m-%d %R:%E9S %Z', field, 'UTC'))", ) self.assertIsNone(params) + + def test_adapt_json_value(self): + import uuid + + from django.core.serializers.json import DjangoJSONEncoder + from django.db.models import JSONField + from google.cloud.spanner_v1 import JsonObject + + self.assertIsNone(self.db_operations.adapt_json_value(None, None)) + obj = JsonObject({"a": 1}) + self.assertEqual(self.db_operations.adapt_json_value(obj, None), {"a": 1}) + adapted_dict = self.db_operations.adapt_json_value({"a": 1}, None) + self.assertIsInstance(adapted_dict, JsonObject) + self.assertEqual(adapted_dict, {"a": 1}) + adapted_list = self.db_operations.adapt_json_value([1, 2], None) + self.assertIsInstance(adapted_list, JsonObject) + self.assertEqual(adapted_list, [1, 2]) + adapted_str = self.db_operations.adapt_json_value("str", None) + self.assertIsInstance(adapted_str, JsonObject) + self.assertEqual(adapted_str, "str") + + u = uuid.UUID("12345678-1234-5678-1234-567812345678") + adapted_uuid = self.db_operations.adapt_json_value({"u": u}, DjangoJSONEncoder) + self.assertIsInstance(adapted_uuid, JsonObject) + self.assertEqual(adapted_uuid, {"u": "12345678-1234-5678-1234-567812345678"}) + + field = JSONField() + self.assertEqual(field.from_db_value(obj, None, None), {"a": 1}) + self.assertEqual(field.from_db_value(JsonObject([1, 2]), None, None), [1, 2]) + self.assertEqual(field.from_db_value(JsonObject("str"), None, None), "str") + self.assertIsNone(field.from_db_value(JsonObject(), None, None)) + self.assertEqual(field.from_db_value({"a": 1}, None, None), {"a": 1}) + self.assertEqual(field.from_db_value('{"a": 1}', None, None), {"a": 1}) + self.assertIsNone(field.from_db_value(None, None, None))