diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py index 553853630..d75037ad1 100644 --- a/google/cloud/bigquery/enums.py +++ b/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): diff --git a/google/cloud/bigquery/standard_sql.py b/google/cloud/bigquery/standard_sql.py index e0f22b2de..68332eb80 100644 --- a/google/cloud/bigquery/standard_sql.py +++ b/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): diff --git a/tests/system/test_client.py b/tests/system/test_client.py index 74c152cf2..04740de8a 100644 --- a/tests/system/test_client.py +++ b/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 '[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, diff --git a/tests/unit/test_standard_sql_types.py b/tests/unit/test_standard_sql_types.py index 0ba0e0cfd..3ed912b5a 100644 --- a/tests/unit/test_standard_sql_types.py +++ b/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):