feat: support range sql

googleapis · Linchin · Mar 7, 2024 · Jan 30, 2024 · Feb 27, 2024 · Mar 5, 2024
commit 0307017eaa64277140072623c567f86402c72686
diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py
@@ -2860,7 +2860,34 @@ def load_table_from_json(
 
  data_str = "\n".join(json.dumps(item, ensure_ascii=False) for item in json_rows)
  encoded_str = data_str.encode()
+ #print(encoded_str)
+ #print(len(encoded_str))
  data_file = io.BytesIO(encoded_str)
+ #print(data_file)
+ #print(data_file.getvalue())
+
+ #with open("temp.json", "w") as f:
+ #b = bytes(encoded_str)
+ #json_file = json.dump(b, f)
+ # f.write(data_str)
+ #f.close()
+
+ if False:
+ with open("temp.json", "rb") as f: 
+ return self.load_table_from_file(
+ f,
+ # data_file,
+ destination,
+ size=len(encoded_str),
+ num_retries=num_retries,
+ job_id=job_id,
+ job_id_prefix=job_id_prefix,
+ location=location,
+ project=project,
+ job_config=new_job_config,
+ timeout=timeout,
+ )
+
  return self.load_table_from_file(
  data_file,
  destination,
@@ -2873,7 +2900,6 @@ def load_table_from_json(
  job_config=new_job_config,
  timeout=timeout,
  )
-
  def _do_resumable_upload(
  self,
  stream: IO[bytes],

@@ -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"
 
 
 class WriteDisposition(object):

@@ -52,6 +52,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 = "RANGE".
  """
 
  def __init__(
@@ -61,12 +63,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 +131,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
+
+ @struct_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 +181,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 +197,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):

@@ -276,6 +276,22 @@ def test_from_api_repr_struct_type_incomplete_field_info(self):
  )
  assert result == expected
 
+ def test_from_api_repr_range_type():
+ pass
+
+ def test_from_api_repr_range_type_missing_element():
+ pass
+
+ def test_from_api_repr_range_type_invalid_element():
+ pass
+
+ def test_to_api_repr_range_type_element_type_missing():
+ pass
+
+ def test_to_api_repr_range_type_w_element_type():
+ pass
+
+
  def test__eq__another_type(self):
  instance = self._make_one()
 
@@ -321,6 +337,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.DATE),
+ ),
  ),
  )
  def test__eq__attribute_differs(self, attr_name, value, value2):