feat: support JSON object consisting of an array. (#782)

Ilya Gurov · web-flow · commit 92a3169b59ba · 2022-08-17T07:50:13.000Z
diff --git a/google/cloud/spanner_v1/data_types.py b/google/cloud/spanner_v1/data_types.py
@@ -19,16 +19,30 @@
 
 class JsonObject(dict):
     """
-    JsonObject type help format Django JSONField to compatible Cloud Spanner's
-    JSON type. Before making queries, it'll help differentiate between
-    normal parameters and JSON parameters.
+    Provides functionality of JSON data type in Cloud Spanner
+    API, mimicking simple `dict()` behaviour and making
+    all the necessary conversions under the hood.
     """
 
     def __init__(self, *args, **kwargs):
         self._is_null = (args, kwargs) == ((), {}) or args == (None,)
+        self._is_array = len(args) and isinstance(args[0], (list, tuple))
+
+        # if the JSON object is represented with an array,
+        # the value is contained separately
+        if self._is_array:
+            self._array_value = args[0]
+            return
+
         if not self._is_null:
             super(JsonObject, self).__init__(*args, **kwargs)
 
+    def __repr__(self):
+        if self._is_array:
+            return str(self._array_value)
+
+        return super(JsonObject, self).__repr__()
+
     @classmethod
     def from_str(cls, str_repr):
         """Initiate an object from its `str` representation.
@@ -53,4 +67,7 @@ def serialize(self):
         if self._is_null:
             return None
 
+        if self._is_array:
+            return json.dumps(self._array_value, sort_keys=True, separators=(",", ":"))
+
         return json.dumps(self, sort_keys=True, separators=(",", ":"))
diff --git a/tests/system/test_dbapi.py b/tests/system/test_dbapi.py
@@ -339,8 +339,10 @@ def test_DDL_autocommit(shared_instance, dbapi_database):
 
 @pytest.mark.skipif(_helpers.USE_EMULATOR, reason="Emulator does not support json.")
 def test_autocommit_with_json_data(shared_instance, dbapi_database):
-    """Check that DDLs in autocommit mode are immediately executed for
-    json fields."""
+    """
+    Check that DDLs in autocommit mode are immediately
+    executed for json fields.
+    """
     # Create table
     conn = Connection(shared_instance, dbapi_database)
     conn.autocommit = True
@@ -376,6 +378,35 @@ def test_autocommit_with_json_data(shared_instance, dbapi_database):
     conn.close()
 
 
+@pytest.mark.skipif(_helpers.USE_EMULATOR, reason="Emulator does not support json.")
+def test_json_array(shared_instance, dbapi_database):
+    # Create table
+    conn = Connection(shared_instance, dbapi_database)
+    conn.autocommit = True
+
+    cur = conn.cursor()
+    cur.execute(
+        """
+        CREATE TABLE JsonDetails (
+            DataId     INT64 NOT NULL,
+            Details    JSON,
+        ) PRIMARY KEY (DataId)
+    """
+    )
+    cur.execute(
+        "INSERT INTO JsonDetails (DataId, Details) VALUES (%s, %s)",
+        [1, JsonObject([1, 2, 3])],
+    )
+
+    cur.execute("SELECT * FROM JsonDetails WHERE DataId = 1")
+    row = cur.fetchone()
+    assert isinstance(row[1], JsonObject)
+    assert row[1].serialize() == "[1,2,3]"
+
+    cur.execute("DROP TABLE JsonDetails")
+    conn.close()
+
+
 def test_DDL_commit(shared_instance, dbapi_database):
     """Check that DDLs in commit mode are executed on calling `commit()`."""
     conn = Connection(shared_instance, dbapi_database)
diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py
@@ -120,6 +120,7 @@
     AllTypesRowData(pkey=108, timestamp_value=NANO_TIME),
     AllTypesRowData(pkey=109, numeric_value=NUMERIC_1),
     AllTypesRowData(pkey=110, json_value=JSON_1),
+    AllTypesRowData(pkey=111, json_value=[JSON_1, JSON_2]),
     # empty array values
     AllTypesRowData(pkey=201, int_array=[]),
     AllTypesRowData(pkey=202, bool_array=[]),
