googleapis · Linchin · Mar 7, 2024 · Jan 30, 2024 · Feb 27, 2024 · Mar 5, 2024
@@ -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):

@@ -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):

@@ -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,

@@ -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):