feat: support insertAll for range (#1909)

Linchin · web-flow · commit 74e75e89ce3a · 2024-05-06T16:46:27.000-07:00
* feat: support insertAll for range

* revert INTERVAL regex

* lint

* add unit test

* lint
diff --git a/google/cloud/bigquery/_helpers.py b/google/cloud/bigquery/_helpers.py
@@ -50,6 +50,7 @@
     r"(?P<days>-?\d+) "
     r"(?P<time_sign>-?)(?P<hours>\d+):(?P<minutes>\d+):(?P<seconds>\d+)\.?(?P<fraction>\d*)?$"
 )
+_RANGE_PATTERN = re.compile(r"\[.*, .*\)")
 
 BIGQUERY_EMULATOR_HOST = "BIGQUERY_EMULATOR_HOST"
 """Environment variable defining host for emulator."""
@@ -334,9 +335,8 @@ def _range_from_json(value, field):
             The parsed range object from ``value`` if the ``field`` is not
             null (otherwise it is :data:`None`).
     """
-    range_literal = re.compile(r"\[.*, .*\)")
     if _not_null(value, field):
-        if range_literal.match(value):
+        if _RANGE_PATTERN.match(value):
             start, end = value[1:-1].split(", ")
             start = _range_element_from_json(start, field.range_element_type)
             end = _range_element_from_json(end, field.range_element_type)
@@ -531,6 +531,52 @@ def _time_to_json(value):
     return value
 
 
+def _range_element_to_json(value, element_type=None):
+    """Coerce 'value' to an JSON-compatible representation."""
+    if value is None:
+        return None
+    elif isinstance(value, str):
+        if value.upper() in ("UNBOUNDED", "NULL"):
+            return None
+        else:
+            # We do not enforce range element value to be valid to reduce
+            # redundancy with backend.
+            return value
+    elif (
+        element_type and element_type.element_type.upper() in _SUPPORTED_RANGE_ELEMENTS
+    ):
+        converter = _SCALAR_VALUE_TO_JSON_ROW.get(element_type.element_type.upper())
+        return converter(value)
+    else:
+        raise ValueError(
+            f"Unsupported RANGE element type {element_type}, or "
+            "element type is empty. Must be DATE, DATETIME, or "
+            "TIMESTAMP"
+        )
+
+
+def _range_field_to_json(range_element_type, value):
+    """Coerce 'value' to an JSON-compatible representation."""
+    if isinstance(value, str):
+        # string literal
+        if _RANGE_PATTERN.match(value):
+            start, end = value[1:-1].split(", ")
+        else:
+            raise ValueError(f"RANGE literal {value} has incorrect format")
+    elif isinstance(value, dict):
+        # dictionary
+        start = value.get("start")
+        end = value.get("end")
+    else:
+        raise ValueError(
+            f"Unsupported type of RANGE value {value}, must be " "string or dict"
+        )
+
+    start = _range_element_to_json(start, range_element_type)
+    end = _range_element_to_json(end, range_element_type)
+    return {"start": start, "end": end}
+
+
 # Converters used for scalar values marshalled to the BigQuery API, such as in
 # query parameters or the tabledata.insert API.
 _SCALAR_VALUE_TO_JSON_ROW = {
@@ -676,6 +722,8 @@ def _single_field_to_json(field, row_value):
 
     if field.field_type == "RECORD":
         return _record_field_to_json(field.fields, row_value)
+    if field.field_type == "RANGE":
+        return _range_field_to_json(field.range_element_type, row_value)
 
     return _scalar_field_to_json(field, row_value)
 
diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py
@@ -1049,10 +1049,22 @@ def test_w_datetime(self):
         self.assertEqual(self._call_fut(when), "12:13:41")
 
 
-def _make_field(field_type, mode="NULLABLE", name="testing", fields=()):
+def _make_field(
+    field_type,
+    mode="NULLABLE",
+    name="testing",
+    fields=(),
+    range_element_type=None,
+):
     from google.cloud.bigquery.schema import SchemaField
 
-    return SchemaField(name=name, field_type=field_type, mode=mode, fields=fields)
+    return SchemaField(
+        name=name,
+        field_type=field_type,
+        mode=mode,
+        fields=fields,
+        range_element_type=range_element_type,
+    )
 
 
 class Test_scalar_field_to_json(unittest.TestCase):
@@ -1251,6 +1263,98 @@ def test_w_dict_unknown_fields(self):
         )
 
 
+class Test_range_field_to_json(unittest.TestCase):
+    def _call_fut(self, field, value):
+        from google.cloud.bigquery._helpers import _range_field_to_json
+
+        return _range_field_to_json(field, value)
+
+    def test_w_date(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        start = datetime.date(2016, 12, 3)
+        original = {"start": start}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_date_string(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        original = {"start": "2016-12-03"}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_datetime(self):
+        field = _make_field("RANGE", range_element_type="DATETIME")
+        start = datetime.datetime(2016, 12, 3, 14, 11, 27, 123456)
+        original = {"start": start}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03T14:11:27.123456", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_datetime_string(self):
+        field = _make_field("RANGE", range_element_type="DATETIME")
+        original = {"start": "2016-12-03T14:11:27.123456"}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03T14:11:27.123456", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_timestamp(self):
+        from google.cloud._helpers import UTC
+
+        field = _make_field("RANGE", range_element_type="TIMESTAMP")
+        start = datetime.datetime(2016, 12, 3, 14, 11, 27, 123456, tzinfo=UTC)
+        original = {"start": start}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03T14:11:27.123456Z", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_timestamp_string(self):
+        field = _make_field("RANGE", range_element_type="TIMESTAMP")
+        original = {"start": "2016-12-03T14:11:27.123456Z"}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03T14:11:27.123456Z", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_timestamp_float(self):
+        field = _make_field("RANGE", range_element_type="TIMESTAMP")
+        original = {"start": 12.34567}
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": 12.34567, "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_string_literal(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        original = "[2016-12-03, UNBOUNDED)"
+        converted = self._call_fut(field.range_element_type, original)
+        expected = {"start": "2016-12-03", "end": None}
+        self.assertEqual(converted, expected)
+
+    def test_w_unsupported_range_element_type(self):
+        field = _make_field("RANGE", range_element_type="TIME")
+        with self.assertRaises(ValueError):
+            self._call_fut(
+                field.range_element_type,
+                {"start": datetime.time(12, 13, 41)},
+            )
+
+    def test_w_no_range_element_type(self):
+        field = _make_field("RANGE")
+        with self.assertRaises(ValueError):
+            self._call_fut(field.range_element_type, "2016-12-03")
+
+    def test_w_incorrect_literal_format(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        original = "[2016-12-03, UNBOUNDED]"
+        with self.assertRaises(ValueError):
+            self._call_fut(field.range_element_type, original)
+
+    def test_w_unsupported_representation(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        with self.assertRaises(ValueError):
+            self._call_fut(field.range_element_type, object())
+
+
 class Test_field_to_json(unittest.TestCase):
     def _call_fut(self, field, value):
         from google.cloud.bigquery._helpers import _field_to_json
@@ -1285,6 +1389,12 @@ def test_w_scalar(self):
         converted = self._call_fut(field, original)
         self.assertEqual(converted, str(original))
 
+    def test_w_range(self):
+        field = _make_field("RANGE", range_element_type="DATE")
+        original = {"start": "2016-12-03", "end": "2024-12-03"}
+        converted = self._call_fut(field, original)
+        self.assertEqual(converted, original)
+
 
 class Test_snake_to_camel_case(unittest.TestCase):
     def _call_fut(self, value):
