feat: support range sql (#1807)

* feat: support range sql

* add unit tests

* add system test

* lint and remove debug code

* lint and remove debug code

* remove added blank line

* add comment for legacy type

Author: Linchin

google/cloud/bigquery/enums.py

@@ -254,6 +254,7 @@ def _generate_next_value_(name, start, count, last_values):
     JSON = enum.auto()
     ARRAY = enum.auto()
     STRUCT = enum.auto()
+    RANGE = enum.auto()
 
 
 class EntityTypes(str, enum.Enum):
@@ -292,6 +293,7 @@ class SqlTypeNames(str, enum.Enum):
     TIME = "TIME"
     DATETIME = "DATETIME"
     INTERVAL = "INTERVAL"  # NOTE: not available in legacy types
+    RANGE = "RANGE"  # NOTE: not available in legacy types
 
 
 class WriteDisposition(object):

google/cloud/bigquery/standard_sql.py

@@ -43,6 +43,7 @@ class StandardSqlDataType:
                 ]
             }
         }
+        RANGE: {type_kind="RANGE", range_element_type="DATETIME"}
 
     Args:
         type_kind:
@@ -52,6 +53,8 @@ class StandardSqlDataType:
             The type of the array's elements, if type_kind is ARRAY.
         struct_type:
             The fields of this struct, in order, if type_kind is STRUCT.
+        range_element_type:
+            The type of the range's elements, if type_kind is RANGE.
     """
 
     def __init__(
@@ -61,12 +64,14 @@ def __init__(
         ] = StandardSqlTypeNames.TYPE_KIND_UNSPECIFIED,
         array_element_type: Optional["StandardSqlDataType"] = None,
         struct_type: Optional["StandardSqlStructType"] = None,
+        range_element_type: Optional["StandardSqlDataType"] = None,
     ):
         self._properties: Dict[str, Any] = {}
 
         self.type_kind = type_kind
         self.array_element_type = array_element_type
         self.struct_type = struct_type
+        self.range_element_type = range_element_type
 
     @property
     def type_kind(self) -> Optional[StandardSqlTypeNames]:
@@ -127,6 +132,28 @@ def struct_type(self, value: Optional["StandardSqlStructType"]):
         else:
             self._properties["structType"] = struct_type
 
+    @property
+    def range_element_type(self) -> Optional["StandardSqlDataType"]:
+        """The type of the range's elements, if type_kind = "RANGE". Must be
+        one of DATETIME, DATE, or TIMESTAMP."""
+        range_element_info = self._properties.get("rangeElementType")
+
+        if range_element_info is None:
+            return None
+
+        result = StandardSqlDataType()
+        result._properties = range_element_info  # We do not use a copy on purpose.
+        return result
+
+    @range_element_type.setter
+    def range_element_type(self, value: Optional["StandardSqlDataType"]):
+        range_element_type = None if value is None else value.to_api_repr()
+
+        if range_element_type is None:
+            self._properties.pop("rangeElementType", None)
+        else:
+            self._properties["rangeElementType"] = range_element_type
+
     def to_api_repr(self) -> Dict[str, Any]:
         """Construct the API resource representation of this SQL data type."""
         return copy.deepcopy(self._properties)
@@ -155,7 +182,13 @@ def from_api_repr(cls, resource: Dict[str, Any]):
             if struct_info:
                 struct_type = StandardSqlStructType.from_api_repr(struct_info)
 
-        return cls(type_kind, array_element_type, struct_type)
+        range_element_type = None
+        if type_kind == StandardSqlTypeNames.RANGE:
+            range_element_info = resource.get("rangeElementType")
+            if range_element_info:
+                range_element_type = cls.from_api_repr(range_element_info)
+
+        return cls(type_kind, array_element_type, struct_type, range_element_type)
 
     def __eq__(self, other):
         if not isinstance(other, StandardSqlDataType):
@@ -165,6 +198,7 @@ def __eq__(self, other):
                 self.type_kind == other.type_kind
                 and self.array_element_type == other.array_element_type
                 and self.struct_type == other.struct_type
+                and self.range_element_type == other.range_element_type
             )
 
     def __str__(self):

tests/system/test_client.py

@@ -2193,6 +2193,44 @@ def test_create_routine(self):
         assert len(rows) == 1
         assert rows[0].max_value == 100.0
 
+    def test_create_routine_with_range(self):
+        routine_name = "routine_range"
+        dataset = self.temp_dataset(_make_dataset_id("routine_range"))
+
+        routine = bigquery.Routine(
+            dataset.routine(routine_name),
+            type_="SCALAR_FUNCTION",
+            language="SQL",
+            body="RANGE_START(x)",
+            arguments=[
+                bigquery.RoutineArgument(
+                    name="x",
+                    data_type=bigquery.StandardSqlDataType(
+                        type_kind=bigquery.StandardSqlTypeNames.RANGE,
+                        range_element_type=bigquery.StandardSqlDataType(
+                            type_kind=bigquery.StandardSqlTypeNames.DATE
+                        ),
+                    ),
+                )
+            ],
+            return_type=bigquery.StandardSqlDataType(
+                type_kind=bigquery.StandardSqlTypeNames.DATE
+            ),
+        )
+
+        query_string = (
+            "SELECT `{}`(RANGE<DATE> '[2016-08-12, UNBOUNDED)') as range_start;".format(
+                str(routine.reference)
+            )
+        )
+
+        routine = helpers.retry_403(Config.CLIENT.create_routine)(routine)
+        query_job = helpers.retry_403(Config.CLIENT.query)(query_string)
+        rows = list(query_job.result())
+
+        assert len(rows) == 1
+        assert rows[0].range_start == datetime.date(2016, 8, 12)
+
     def test_create_tvf_routine(self):
         from google.cloud.bigquery import (
             Routine,

tests/unit/test_standard_sql_types.py

@@ -129,6 +129,28 @@ def test_to_api_repr_struct_type_w_field_types(self):
         }
         assert result == expected
 
+    def test_to_api_repr_range_type_element_type_missing(self):
+        instance = self._make_one(
+            bq.StandardSqlTypeNames.RANGE, range_element_type=None
+        )
+
+        result = instance.to_api_repr()
+
+        assert result == {"typeKind": "RANGE"}
+
+    def test_to_api_repr_range_type_w_element_type(self):
+        range_element_type = self._make_one(type_kind=bq.StandardSqlTypeNames.DATE)
+        instance = self._make_one(
+            bq.StandardSqlTypeNames.RANGE, range_element_type=range_element_type
+        )
+
+        result = instance.to_api_repr()
+
+        assert result == {
+            "typeKind": "RANGE",
+            "rangeElementType": {"typeKind": "DATE"},
+        }
+
     def test_from_api_repr_empty_resource(self):
         klass = self._get_target_class()
         result = klass.from_api_repr(resource={})
@@ -276,6 +298,31 @@ def test_from_api_repr_struct_type_incomplete_field_info(self):
         )
         assert result == expected
 
+    def test_from_api_repr_range_type_full(self):
+        klass = self._get_target_class()
+        resource = {"typeKind": "RANGE", "rangeElementType": {"typeKind": "DATE"}}
+
+        result = klass.from_api_repr(resource=resource)
+
+        expected = klass(
+            type_kind=bq.StandardSqlTypeNames.RANGE,
+            range_element_type=klass(type_kind=bq.StandardSqlTypeNames.DATE),
+        )
+        assert result == expected
+
+    def test_from_api_repr_range_type_missing_element_type(self):
+        klass = self._get_target_class()
+        resource = {"typeKind": "RANGE"}
+
+        result = klass.from_api_repr(resource=resource)
+
+        expected = klass(
+            type_kind=bq.StandardSqlTypeNames.RANGE,
+            range_element_type=None,
+            struct_type=None,
+        )
+        assert result == expected
+
     def test__eq__another_type(self):
         instance = self._make_one()
 
@@ -321,6 +368,11 @@ def test__eq__similar_instance(self):
                 bq.StandardSqlStructType(fields=[bq.StandardSqlField(name="foo")]),
                 bq.StandardSqlStructType(fields=[bq.StandardSqlField(name="bar")]),
             ),
+            (
+                "range_element_type",
+                bq.StandardSqlDataType(type_kind=bq.StandardSqlTypeNames.DATE),
+                bq.StandardSqlDataType(type_kind=bq.StandardSqlTypeNames.DATETIME),
+            ),
         ),
     )
     def test__eq__attribute_differs(self, attr_name, value, value2):