Skip to content

Commit b359a9a

Browse files
LinchinLinchin
authored
feat: support RANGE query parameters (#1827)
* feat: RANGE query parameters and unit tests * unit test * unit test coverage * lint * lint * lint * system test * fix system test * ajust init items order * fix typos and improve docstrings
1 parent e81a13c commit b359a9a

File tree

5 files changed

+883
-1
lines changed

5 files changed

+883
-1
lines changed

‎benchmark/benchmark.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -231,7 +231,7 @@ def _is_datetime_min(time_str: str) -> bool:
231231

232232

233233
def _summary(run: dict) -> str:
234-
"""Coverts run dict to run summary string."""
234+
"""Converts run dict to run summary string."""
235235
no_val = "NODATA"
236236
output = ["QUERYTIME "]
237237

‎google/cloud/bigquery/__init__.py

+4
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@
8383
from google.cloud.bigquery.query import ConnectionProperty
8484
from google.cloud.bigquery.query import ScalarQueryParameter
8585
from google.cloud.bigquery.query import ScalarQueryParameterType
86+
from google.cloud.bigquery.query import RangeQueryParameter
87+
from google.cloud.bigquery.query import RangeQueryParameterType
8688
from google.cloud.bigquery.query import SqlParameterScalarTypes
8789
from google.cloud.bigquery.query import StructQueryParameter
8890
from google.cloud.bigquery.query import StructQueryParameterType
@@ -122,10 +124,12 @@
122124
"ArrayQueryParameter",
123125
"ScalarQueryParameter",
124126
"StructQueryParameter",
127+
"RangeQueryParameter",
125128
"ArrayQueryParameterType",
126129
"ScalarQueryParameterType",
127130
"SqlParameterScalarTypes",
128131
"StructQueryParameterType",
132+
"RangeQueryParameterType",
129133
# Datasets
130134
"Dataset",
131135
"DatasetReference",

‎google/cloud/bigquery/query.py

+297
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@
3030
Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]
3131
]
3232

33+
_RANGE_ELEMENT_TYPE_STR = {"TIMESTAMP", "DATETIME", "DATE"}
34+
3335

3436
class ConnectionProperty:
3537
"""A connection-level property to customize query behavior.
@@ -362,6 +364,129 @@ def __repr__(self):
362364
return f"{self.__class__.__name__}({items}{name}{description})"
363365

364366

367+
class RangeQueryParameterType(_AbstractQueryParameterType):
368+
"""Type representation for range query parameters.
369+
370+
Args:
371+
type_ (Union[ScalarQueryParameterType, str]):
372+
Type of range element, must be one of 'TIMESTAMP', 'DATETIME', or
373+
'DATE'.
374+
name (Optional[str]):
375+
The name of the query parameter. Primarily used if the type is
376+
one of the subfields in ``StructQueryParameterType`` instance.
377+
description (Optional[str]):
378+
The query parameter description. Primarily used if the type is
379+
one of the subfields in ``StructQueryParameterType`` instance.
380+
"""
381+
382+
@classmethod
383+
def _parse_range_element_type(self, type_):
384+
"""Helper method that parses the input range element type, which may
385+
be a string, or a ScalarQueryParameterType object.
386+
387+
Returns:
388+
google.cloud.bigquery.query.ScalarQueryParameterType: Instance
389+
"""
390+
if isinstance(type_, str):
391+
if type_ not in _RANGE_ELEMENT_TYPE_STR:
392+
raise ValueError(
393+
"If given as a string, range element type must be one of "
394+
"'TIMESTAMP', 'DATE', or 'DATETIME'."
395+
)
396+
return ScalarQueryParameterType(type_)
397+
elif isinstance(type_, ScalarQueryParameterType):
398+
if type_._type not in _RANGE_ELEMENT_TYPE_STR:
399+
raise ValueError(
400+
"If given as a ScalarQueryParameter object, range element "
401+
"type must be one of 'TIMESTAMP', 'DATE', or 'DATETIME' "
402+
"type."
403+
)
404+
return type_
405+
else:
406+
raise ValueError(
407+
"range_type must be a string or ScalarQueryParameter object, "
408+
"of 'TIMESTAMP', 'DATE', or 'DATETIME' type."
409+
)
410+
411+
def __init__(self, type_, *, name=None, description=None):
412+
self.type_ = self._parse_range_element_type(type_)
413+
self.name = name
414+
self.description = description
415+
416+
@classmethod
417+
def from_api_repr(cls, resource):
418+
"""Factory: construct parameter type from JSON resource.
419+
420+
Args:
421+
resource (Dict): JSON mapping of parameter
422+
423+
Returns:
424+
google.cloud.bigquery.query.RangeQueryParameterType: Instance
425+
"""
426+
type_ = resource["rangeElementType"]["type"]
427+
name = resource.get("name")
428+
description = resource.get("description")
429+
430+
return cls(type_, name=name, description=description)
431+
432+
def to_api_repr(self):
433+
"""Construct JSON API representation for the parameter type.
434+
435+
Returns:
436+
Dict: JSON mapping
437+
"""
438+
# Name and description are only used if the type is a field inside a struct
439+
# type, but it's StructQueryParameterType's responsibilty to use these two
440+
# attributes in the API representation when needed. Here we omit them.
441+
return {
442+
"type": "RANGE",
443+
"rangeElementType": self.type_.to_api_repr(),
444+
}
445+
446+
def with_name(self, new_name: Union[str, None]):
447+
"""Return a copy of the instance with ``name`` set to ``new_name``.
448+
449+
Args:
450+
name (Union[str, None]):
451+
The new name of the range query parameter type. If ``None``,
452+
the existing name is cleared.
453+
454+
Returns:
455+
google.cloud.bigquery.query.RangeQueryParameterType:
456+
A new instance with updated name.
457+
"""
458+
return type(self)(self.type_, name=new_name, description=self.description)
459+
460+
def __repr__(self):
461+
name = f", name={self.name!r}" if self.name is not None else ""
462+
description = (
463+
f", description={self.description!r}"
464+
if self.description is not None
465+
else ""
466+
)
467+
return f"{self.__class__.__name__}({self.type_!r}{name}{description})"
468+
469+
def _key(self):
470+
"""A tuple key that uniquely describes this field.
471+
472+
Used to compute this instance's hashcode and evaluate equality.
473+
474+
Returns:
475+
Tuple: The contents of this
476+
:class:`~google.cloud.bigquery.query.RangeQueryParameterType`.
477+
"""
478+
type_ = self.type_.to_api_repr()
479+
return (self.name, type_, self.description)
480+
481+
def __eq__(self, other):
482+
if not isinstance(other, RangeQueryParameterType):
483+
return NotImplemented
484+
return self._key() == other._key()
485+
486+
def __ne__(self, other):
487+
return not self == other
488+
489+
365490
class _AbstractQueryParameter(object):
366491
"""Base class for named / positional query parameters."""
367492

@@ -811,6 +936,178 @@ def __repr__(self):
811936
return "StructQueryParameter{}".format(self._key())
812937

813938

939+
class RangeQueryParameter(_AbstractQueryParameter):
940+
"""Named / positional query parameters for range values.
941+
942+
Args:
943+
range_element_type (Union[str, RangeQueryParameterType]):
944+
The type of range elements. It must be one of 'TIMESTAMP',
945+
'DATE', or 'DATETIME'.
946+
947+
start (Optional[Union[ScalarQueryParameter, str]]):
948+
The start of the range value. Must be the same type as
949+
range_element_type. If not provided, it's interpreted as UNBOUNDED.
950+
951+
end (Optional[Union[ScalarQueryParameter, str]]):
952+
The end of the range value. Must be the same type as
953+
range_element_type. If not provided, it's interpreted as UNBOUNDED.
954+
955+
name (Optional[str]):
956+
Parameter name, used via ``@foo`` syntax. If None, the
957+
parameter can only be addressed via position (``?``).
958+
"""
959+
960+
@classmethod
961+
def _parse_range_element_type(self, range_element_type):
962+
if isinstance(range_element_type, str):
963+
if range_element_type not in _RANGE_ELEMENT_TYPE_STR:
964+
raise ValueError(
965+
"If given as a string, range_element_type must be one of "
966+
f"'TIMESTAMP', 'DATE', or 'DATETIME'. Got {range_element_type}."
967+
)
968+
return RangeQueryParameterType(range_element_type)
969+
elif isinstance(range_element_type, RangeQueryParameterType):
970+
if range_element_type.type_._type not in _RANGE_ELEMENT_TYPE_STR:
971+
raise ValueError(
972+
"If given as a RangeQueryParameterType object, "
973+
"range_element_type must be one of 'TIMESTAMP', 'DATE', "
974+
"or 'DATETIME' type."
975+
)
976+
return range_element_type
977+
else:
978+
raise ValueError(
979+
"range_element_type must be a string or "
980+
"RangeQueryParameterType object, of 'TIMESTAMP', 'DATE', "
981+
"or 'DATETIME' type. Got "
982+
f"{type(range_element_type)}:{range_element_type}"
983+
)
984+
985+
@classmethod
986+
def _serialize_range_element_value(self, value, type_):
987+
if value is None or isinstance(value, str):
988+
return value
989+
else:
990+
converter = _SCALAR_VALUE_TO_JSON_PARAM.get(type_)
991+
if converter is not None:
992+
return converter(value) # type: ignore
993+
else:
994+
raise ValueError(
995+
f"Cannot convert range element value from type {type_}, "
996+
"must be one of the strings 'TIMESTAMP', 'DATE' "
997+
"'DATETIME' or a RangeQueryParameterType object."
998+
)
999+
1000+
def __init__(
1001+
self,
1002+
range_element_type,
1003+
start=None,
1004+
end=None,
1005+
name=None,
1006+
):
1007+
self.name = name
1008+
self.range_element_type = self._parse_range_element_type(range_element_type)
1009+
print(self.range_element_type.type_._type)
1010+
self.start = start
1011+
self.end = end
1012+
1013+
@classmethod
1014+
def positional(
1015+
cls, range_element_type, start=None, end=None
1016+
) -> "RangeQueryParameter":
1017+
"""Factory for positional parameters.
1018+
1019+
Args:
1020+
range_element_type (Union[str, RangeQueryParameterType]):
1021+
The type of range elements. It must be one of `'TIMESTAMP'`,
1022+
`'DATE'`, or `'DATETIME'`.
1023+
1024+
start (Optional[Union[ScalarQueryParameter, str]]):
1025+
The start of the range value. Must be the same type as
1026+
range_element_type. If not provided, it's interpreted as
1027+
UNBOUNDED.
1028+
1029+
end (Optional[Union[ScalarQueryParameter, str]]):
1030+
The end of the range value. Must be the same type as
1031+
range_element_type. If not provided, it's interpreted as
1032+
UNBOUNDED.
1033+
1034+
Returns:
1035+
google.cloud.bigquery.query.RangeQueryParameter: Instance without
1036+
name.
1037+
"""
1038+
return cls(range_element_type, start, end)
1039+
1040+
@classmethod
1041+
def from_api_repr(cls, resource: dict) -> "RangeQueryParameter":
1042+
"""Factory: construct parameter from JSON resource.
1043+
1044+
Args:
1045+
resource (Dict): JSON mapping of parameter
1046+
1047+
Returns:
1048+
google.cloud.bigquery.query.RangeQueryParameter: Instance
1049+
"""
1050+
name = resource.get("name")
1051+
range_element_type = (
1052+
resource.get("parameterType", {}).get("rangeElementType", {}).get("type")
1053+
)
1054+
range_value = resource.get("parameterValue", {}).get("rangeValue", {})
1055+
start = range_value.get("start", {}).get("value")
1056+
end = range_value.get("end", {}).get("value")
1057+
1058+
return cls(range_element_type, start=start, end=end, name=name)
1059+
1060+
def to_api_repr(self) -> dict:
1061+
"""Construct JSON API representation for the parameter.
1062+
1063+
Returns:
1064+
Dict: JSON mapping
1065+
"""
1066+
range_element_type = self.range_element_type.to_api_repr()
1067+
type_ = self.range_element_type.type_._type
1068+
start = self._serialize_range_element_value(self.start, type_)
1069+
end = self._serialize_range_element_value(self.end, type_)
1070+
resource = {
1071+
"parameterType": range_element_type,
1072+
"parameterValue": {
1073+
"rangeValue": {
1074+
"start": {"value": start},
1075+
"end": {"value": end},
1076+
},
1077+
},
1078+
}
1079+
1080+
# distinguish between name not provided vs. name being empty string
1081+
if self.name is not None:
1082+
resource["name"] = self.name
1083+
1084+
return resource
1085+
1086+
def _key(self):
1087+
"""A tuple key that uniquely describes this field.
1088+
1089+
Used to compute this instance's hashcode and evaluate equality.
1090+
1091+
Returns:
1092+
Tuple: The contents of this
1093+
:class:`~google.cloud.bigquery.query.RangeQueryParameter`.
1094+
"""
1095+
1096+
range_element_type = self.range_element_type.to_api_repr()
1097+
return (self.name, range_element_type, self.start, self.end)
1098+
1099+
def __eq__(self, other):
1100+
if not isinstance(other, RangeQueryParameter):
1101+
return NotImplemented
1102+
return self._key() == other._key()
1103+
1104+
def __ne__(self, other):
1105+
return not self == other
1106+
1107+
def __repr__(self):
1108+
return "RangeQueryParameter{}".format(self._key())
1109+
1110+
8141111
class SqlParameterScalarTypes:
8151112
"""Supported scalar SQL query parameter types as type objects."""
8161113

‎tests/system/test_query.py

+33
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from google.cloud.bigquery.query import ScalarQueryParameterType
2727
from google.cloud.bigquery.query import StructQueryParameter
2828
from google.cloud.bigquery.query import StructQueryParameterType
29+
from google.cloud.bigquery.query import RangeQueryParameter
2930

3031

3132
@pytest.fixture(params=["INSERT", "QUERY"])
@@ -422,6 +423,38 @@ def test_query_statistics(bigquery_client, query_api_method):
422423
)
423424
],
424425
),
426+
(
427+
"SELECT @range_date",
428+
"[2016-12-05, UNBOUNDED)",
429+
[
430+
RangeQueryParameter(
431+
name="range_date",
432+
range_element_type="DATE",
433+
start=datetime.date(2016, 12, 5),
434+
)
435+
],
436+
),
437+
(
438+
"SELECT @range_datetime",
439+
"[2016-12-05T00:00:00, UNBOUNDED)",
440+
[
441+
RangeQueryParameter(
442+
name="range_datetime",
443+
range_element_type="DATETIME",
444+
start=datetime.datetime(2016, 12, 5),
445+
)
446+
],
447+
),
448+
(
449+
"SELECT @range_unbounded",
450+
"[UNBOUNDED, UNBOUNDED)",
451+
[
452+
RangeQueryParameter(
453+
name="range_unbounded",
454+
range_element_type="DATETIME",
455+
)
456+
],
457+
),
425458
),
426459
)
427460
def test_query_parameters(

0 commit comments

Comments
 (0)