From 2d19fcc94365d937060804b4e001a2a415e053e7 Mon Sep 17 00:00:00 2001
From: Tim Swast <swast@google.com>
Date: Fri, 15 Dec 2023 16:17:44 -0600
Subject: [PATCH] feat: support JSON type in `insert_rows` and as a scalar
 query parameter

---
 google/cloud/bigquery/_helpers.py |  9 ++++++++-
 tests/system/test_client.py       |  7 ++++++-
 tests/system/test_query.py        | 12 ++++++++++++
 tests/unit/test__helpers.py       | 16 ++++++++++++++++
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
index 93b46341e..4cf6dddac 100644
--- a/google/cloud/bigquery/_helpers.py
+++ b/google/cloud/bigquery/_helpers.py
@@ -374,6 +374,13 @@ def _bytes_to_json(value):
     return value
 
 
+def _json_to_json(value):
+    """Coerce 'value' to a BigQuery REST API representation."""
+    if value is None:
+        return None
+    return json.dumps(value)
+
+
 def _timestamp_to_json_parameter(value):
     """Coerce 'value' to an JSON-compatible representation.
 
@@ -439,7 +446,7 @@ def _time_to_json(value):
     "DATETIME": _datetime_to_json,
     "DATE": _date_to_json,
     "TIME": _time_to_json,
-    "JSON": _json_from_json,
+    "JSON": _json_to_json,
     # Make sure DECIMAL and BIGDECIMAL are handled, even though
     # requests for them should be converted to NUMERIC.  Better safe
     # than sorry.
diff --git a/tests/system/test_client.py b/tests/system/test_client.py
index 7cea8cfa4..92894455a 100644
--- a/tests/system/test_client.py
+++ b/tests/system/test_client.py
@@ -2049,13 +2049,18 @@ def test_insert_rows_nested_nested(self):
                     ),
                 ],
             ),
+            SF("json_col", "JSON"),
         ]
         record = {
             "nested_string": "another string value",
             "nested_repeated": [0, 1, 2],
             "nested_record": {"nested_nested_string": "some deep insight"},
         }
-        to_insert = [("Some value", record)]
+        json_record = {
+            "json_array": [1, 2, 3],
+            "json_object": {"alpha": "abc", "num": 123},
+        }
+        to_insert = [("Some value", record, json_record)]
         table_id = "test_table"
         dataset = self.temp_dataset(_make_dataset_id("issue_2951"))
         table_arg = Table(dataset.table(table_id), schema=schema)
diff --git a/tests/system/test_query.py b/tests/system/test_query.py
index 723f927d7..b8e0c00da 100644
--- a/tests/system/test_query.py
+++ b/tests/system/test_query.py
@@ -256,6 +256,18 @@ def test_query_statistics(bigquery_client, query_api_method):
                 )
             ],
         ),
+        pytest.param(
+            "SELECT @json",
+            {"alpha": "abc", "num": [1, 2, 3]},
+            [
+                ScalarQueryParameter(
+                    name="json",
+                    type_="JSON",
+                    value={"alpha": "abc", "num": [1, 2, 3]},
+                )
+            ],
+            id="scalar-json",
+        ),
         (
             "SELECT @naive_time",
             datetime.time(12, 41, 9, 62500),
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
index 7bf55baeb..87ab46669 100644
--- a/tests/unit/test__helpers.py
+++ b/tests/unit/test__helpers.py
@@ -886,6 +886,16 @@ def test_w_known_field_type(self):
         converted = self._call_fut(field, original)
         self.assertEqual(converted, str(original))
 
+    def test_w_scalar_none(self):
+        import google.cloud.bigquery._helpers as module_under_test
+
+        scalar_types = module_under_test._SCALAR_VALUE_TO_JSON_ROW.keys()
+        for type_ in scalar_types:
+            field = _make_field(type_)
+            original = None
+            converted = self._call_fut(field, original)
+            self.assertIsNone(converted, msg=f"{type_} did not return None")
+
 
 class Test_single_field_to_json(unittest.TestCase):
     def _call_fut(self, field, value):
@@ -921,6 +931,12 @@ def test_w_scalar_ignores_mode(self):
         converted = self._call_fut(field, original)
         self.assertEqual(converted, original)
 
+    def test_w_scalar_json(self):
+        field = _make_field("JSON")
+        original = {"alpha": "abc", "num": [1, 2, 3]}
+        converted = self._call_fut(field, original)
+        self.assertEqual(converted, json.dumps(original))
+
 
 class Test_repeated_field_to_json(unittest.TestCase):
     def _call_fut(self, field, value):