Skip to content

Commit

Permalink
feat: support range sql
Browse files Browse the repository at this point in the history
  • Loading branch information
Linchin committed Jan 30, 2024
1 parent 6249032 commit 0307017
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 2 deletions.
28 changes: 27 additions & 1 deletion google/cloud/bigquery/client.py
Expand Up @@ -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,
Expand All @@ -2873,7 +2900,6 @@ def load_table_from_json(
job_config=new_job_config,
timeout=timeout,
)

def _do_resumable_upload(
self,
stream: IO[bytes],
Expand Down
2 changes: 2 additions & 0 deletions google/cloud/bigquery/enums.py
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
35 changes: 34 additions & 1 deletion google/cloud/bigquery/standard_sql.py
Expand Up @@ -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__(
Expand All @@ -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]:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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):
Expand All @@ -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):
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/test_standard_sql_types.py
Expand Up @@ -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()

Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 0307017

Please sign in to comment.