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
googleapis · Mar 4, 2024 · b359a9a · b359a9a
1 parent e81a13c
commit b359a9a
Show file tree

Hide file tree

Showing 5 changed files with 883 additions and 1 deletion.
diff --git a/benchmark/benchmark.py b/benchmark/benchmark.py
@@ -231,7 +231,7 @@ def _is_datetime_min(time_str: str) -> bool:
 
 
 def _summary(run: dict) -> str:
-    """Coverts run dict to run summary string."""
+    """Converts run dict to run summary string."""
     no_val = "NODATA"
     output = ["QUERYTIME "]
 

diff --git a/google/cloud/bigquery/__init__.py b/google/cloud/bigquery/__init__.py
@@ -83,6 +83,8 @@
 from google.cloud.bigquery.query import ConnectionProperty
 from google.cloud.bigquery.query import ScalarQueryParameter
 from google.cloud.bigquery.query import ScalarQueryParameterType
+from google.cloud.bigquery.query import RangeQueryParameter
+from google.cloud.bigquery.query import RangeQueryParameterType
 from google.cloud.bigquery.query import SqlParameterScalarTypes
 from google.cloud.bigquery.query import StructQueryParameter
 from google.cloud.bigquery.query import StructQueryParameterType
@@ -122,10 +124,12 @@
     "ArrayQueryParameter",
     "ScalarQueryParameter",
     "StructQueryParameter",
+    "RangeQueryParameter",
     "ArrayQueryParameterType",
     "ScalarQueryParameterType",
     "SqlParameterScalarTypes",
     "StructQueryParameterType",
+    "RangeQueryParameterType",
     # Datasets
     "Dataset",
     "DatasetReference",

diff --git a/google/cloud/bigquery/query.py b/google/cloud/bigquery/query.py
@@ -30,6 +30,8 @@
     Union[str, int, float, decimal.Decimal, bool, datetime.datetime, datetime.date]
 ]
 
+_RANGE_ELEMENT_TYPE_STR = {"TIMESTAMP", "DATETIME", "DATE"}
+
 
 class ConnectionProperty:
     """A connection-level property to customize query behavior.
@@ -362,6 +364,129 @@ def __repr__(self):
         return f"{self.__class__.__name__}({items}{name}{description})"
 
 
+class RangeQueryParameterType(_AbstractQueryParameterType):
+    """Type representation for range query parameters.
+
+    Args:
+        type_ (Union[ScalarQueryParameterType, str]):
+            Type of range element, must be one of 'TIMESTAMP', 'DATETIME', or
+            'DATE'.
+        name (Optional[str]):
+            The name of the query parameter. Primarily used if the type is
+            one of the subfields in ``StructQueryParameterType`` instance.
+        description (Optional[str]):
+            The query parameter description. Primarily used if the type is
+            one of the subfields in ``StructQueryParameterType`` instance.
+    """
+
+    @classmethod
+    def _parse_range_element_type(self, type_):
+        """Helper method that parses the input range element type, which may
+        be a string, or a ScalarQueryParameterType object.
+
+        Returns:
+            google.cloud.bigquery.query.ScalarQueryParameterType: Instance
+        """
+        if isinstance(type_, str):
+            if type_ not in _RANGE_ELEMENT_TYPE_STR:
+                raise ValueError(
+                    "If given as a string, range element type must be one of "
+                    "'TIMESTAMP', 'DATE', or 'DATETIME'."
+                )
+            return ScalarQueryParameterType(type_)
+        elif isinstance(type_, ScalarQueryParameterType):
+            if type_._type not in _RANGE_ELEMENT_TYPE_STR:
+                raise ValueError(
+                    "If given as a ScalarQueryParameter object, range element "
+                    "type must be one of 'TIMESTAMP', 'DATE', or 'DATETIME' "
+                    "type."
+                )
+            return type_
+        else:
+            raise ValueError(
+                "range_type must be a string or ScalarQueryParameter object, "
+                "of 'TIMESTAMP', 'DATE', or 'DATETIME' type."
+            )
+
+    def __init__(self, type_, *, name=None, description=None):
+        self.type_ = self._parse_range_element_type(type_)
+        self.name = name
+        self.description = description
+
+    @classmethod
+    def from_api_repr(cls, resource):
+        """Factory: construct parameter type from JSON resource.
+
+        Args:
+            resource (Dict): JSON mapping of parameter
+
+        Returns:
+            google.cloud.bigquery.query.RangeQueryParameterType: Instance
+        """
+        type_ = resource["rangeElementType"]["type"]
+        name = resource.get("name")
+        description = resource.get("description")
+
+        return cls(type_, name=name, description=description)
+
+    def to_api_repr(self):
+        """Construct JSON API representation for the parameter type.
+
+        Returns:
+            Dict: JSON mapping
+        """
+        # Name and description are only used if the type is a field inside a struct
+        # type, but it's StructQueryParameterType's responsibilty to use these two
+        # attributes in the API representation when needed. Here we omit them.
+        return {
+            "type": "RANGE",
+            "rangeElementType": self.type_.to_api_repr(),
+        }
+
+    def with_name(self, new_name: Union[str, None]):
+        """Return a copy of the instance with ``name`` set to ``new_name``.
+
+        Args:
+            name (Union[str, None]):
+                The new name of the range query parameter type. If ``None``,
+                the existing name is cleared.
+
+        Returns:
+            google.cloud.bigquery.query.RangeQueryParameterType:
+               A new instance with updated name.
+        """
+        return type(self)(self.type_, name=new_name, description=self.description)
+
+    def __repr__(self):
+        name = f", name={self.name!r}" if self.name is not None else ""
+        description = (
+            f", description={self.description!r}"
+            if self.description is not None
+            else ""
+        )
+        return f"{self.__class__.__name__}({self.type_!r}{name}{description})"
+
+    def _key(self):
+        """A tuple key that uniquely describes this field.
+
+        Used to compute this instance's hashcode and evaluate equality.
+
+        Returns:
+            Tuple: The contents of this
+            :class:`~google.cloud.bigquery.query.RangeQueryParameterType`.
+        """
+        type_ = self.type_.to_api_repr()
+        return (self.name, type_, self.description)
+
+    def __eq__(self, other):
+        if not isinstance(other, RangeQueryParameterType):
+            return NotImplemented
+        return self._key() == other._key()
+
+    def __ne__(self, other):
+        return not self == other
+
+
 class _AbstractQueryParameter(object):
     """Base class for named / positional query parameters."""
 
@@ -811,6 +936,178 @@ def __repr__(self):
         return "StructQueryParameter{}".format(self._key())
 
 
+class RangeQueryParameter(_AbstractQueryParameter):
+    """Named / positional query parameters for range values.
+
+    Args:
+        range_element_type (Union[str, RangeQueryParameterType]):
+            The type of range elements. It must be one of 'TIMESTAMP',
+            'DATE', or 'DATETIME'.
+
+        start (Optional[Union[ScalarQueryParameter, str]]):
+            The start of the range value. Must be the same type as
+            range_element_type. If not provided, it's interpreted as UNBOUNDED.
+
+        end (Optional[Union[ScalarQueryParameter, str]]):
+            The end of the range value. Must be the same type as
+            range_element_type. If not provided, it's interpreted as UNBOUNDED.
+
+        name (Optional[str]):
+            Parameter name, used via ``@foo`` syntax.  If None, the
+            parameter can only be addressed via position (``?``).
+    """
+
+    @classmethod
+    def _parse_range_element_type(self, range_element_type):
+        if isinstance(range_element_type, str):
+            if range_element_type not in _RANGE_ELEMENT_TYPE_STR:
+                raise ValueError(
+                    "If given as a string, range_element_type must be one of "
+                    f"'TIMESTAMP', 'DATE', or 'DATETIME'. Got {range_element_type}."
+                )
+            return RangeQueryParameterType(range_element_type)
+        elif isinstance(range_element_type, RangeQueryParameterType):
+            if range_element_type.type_._type not in _RANGE_ELEMENT_TYPE_STR:
+                raise ValueError(
+                    "If given as a RangeQueryParameterType object, "
+                    "range_element_type must be one of 'TIMESTAMP', 'DATE', "
+                    "or 'DATETIME' type."
+                )
+            return range_element_type
+        else:
+            raise ValueError(
+                "range_element_type must be a string or "
+                "RangeQueryParameterType object, of 'TIMESTAMP', 'DATE', "
+                "or 'DATETIME' type. Got "
+                f"{type(range_element_type)}:{range_element_type}"
+            )
+
+    @classmethod
+    def _serialize_range_element_value(self, value, type_):
+        if value is None or isinstance(value, str):
+            return value
+        else:
+            converter = _SCALAR_VALUE_TO_JSON_PARAM.get(type_)
+            if converter is not None:
+                return converter(value)  # type: ignore
+            else:
+                raise ValueError(
+                    f"Cannot convert range element value from type {type_}, "
+                    "must be one of the strings 'TIMESTAMP', 'DATE' "
+                    "'DATETIME' or a RangeQueryParameterType object."
+                )
+
+    def __init__(
+        self,
+        range_element_type,
+        start=None,
+        end=None,
+        name=None,
+    ):
+        self.name = name
+        self.range_element_type = self._parse_range_element_type(range_element_type)
+        print(self.range_element_type.type_._type)
+        self.start = start
+        self.end = end
+
+    @classmethod
+    def positional(
+        cls, range_element_type, start=None, end=None
+    ) -> "RangeQueryParameter":
+        """Factory for positional parameters.
+
+        Args:
+            range_element_type (Union[str, RangeQueryParameterType]):
+                The type of range elements. It must be one of `'TIMESTAMP'`,
+                `'DATE'`, or `'DATETIME'`.
+
+            start (Optional[Union[ScalarQueryParameter, str]]):
+                The start of the range value. Must be the same type as
+                range_element_type. If not provided, it's interpreted as
+                UNBOUNDED.
+
+            end (Optional[Union[ScalarQueryParameter, str]]):
+                The end of the range value. Must be the same type as
+                range_element_type. If not provided, it's interpreted as
+                UNBOUNDED.
+
+        Returns:
+            google.cloud.bigquery.query.RangeQueryParameter: Instance without
+            name.
+        """
+        return cls(range_element_type, start, end)
+
+    @classmethod
+    def from_api_repr(cls, resource: dict) -> "RangeQueryParameter":
+        """Factory: construct parameter from JSON resource.
+
+        Args:
+            resource (Dict): JSON mapping of parameter
+
+        Returns:
+            google.cloud.bigquery.query.RangeQueryParameter: Instance
+        """
+        name = resource.get("name")
+        range_element_type = (
+            resource.get("parameterType", {}).get("rangeElementType", {}).get("type")
+        )
+        range_value = resource.get("parameterValue", {}).get("rangeValue", {})
+        start = range_value.get("start", {}).get("value")
+        end = range_value.get("end", {}).get("value")
+
+        return cls(range_element_type, start=start, end=end, name=name)
+
+    def to_api_repr(self) -> dict:
+        """Construct JSON API representation for the parameter.
+
+        Returns:
+            Dict: JSON mapping
+        """
+        range_element_type = self.range_element_type.to_api_repr()
+        type_ = self.range_element_type.type_._type
+        start = self._serialize_range_element_value(self.start, type_)
+        end = self._serialize_range_element_value(self.end, type_)
+        resource = {
+            "parameterType": range_element_type,
+            "parameterValue": {
+                "rangeValue": {
+                    "start": {"value": start},
+                    "end": {"value": end},
+                },
+            },
+        }
+
+        # distinguish between name not provided vs. name being empty string
+        if self.name is not None:
+            resource["name"] = self.name
+
+        return resource
+
+    def _key(self):
+        """A tuple key that uniquely describes this field.
+
+        Used to compute this instance's hashcode and evaluate equality.
+
+        Returns:
+            Tuple: The contents of this
+            :class:`~google.cloud.bigquery.query.RangeQueryParameter`.
+        """
+
+        range_element_type = self.range_element_type.to_api_repr()
+        return (self.name, range_element_type, self.start, self.end)
+
+    def __eq__(self, other):
+        if not isinstance(other, RangeQueryParameter):
+            return NotImplemented
+        return self._key() == other._key()
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __repr__(self):
+        return "RangeQueryParameter{}".format(self._key())
+
+
 class SqlParameterScalarTypes:
     """Supported scalar SQL query parameter types as type objects."""
 

diff --git a/tests/system/test_query.py b/tests/system/test_query.py
@@ -26,6 +26,7 @@
 from google.cloud.bigquery.query import ScalarQueryParameterType
 from google.cloud.bigquery.query import StructQueryParameter
 from google.cloud.bigquery.query import StructQueryParameterType
+from google.cloud.bigquery.query import RangeQueryParameter
 
 
 @pytest.fixture(params=["INSERT", "QUERY"])
@@ -422,6 +423,38 @@ def test_query_statistics(bigquery_client, query_api_method):
                 )
             ],
         ),
+        (
+            "SELECT @range_date",
+            "[2016-12-05, UNBOUNDED)",
+            [
+                RangeQueryParameter(
+                    name="range_date",
+                    range_element_type="DATE",
+                    start=datetime.date(2016, 12, 5),
+                )
+            ],
+        ),
+        (
+            "SELECT @range_datetime",
+            "[2016-12-05T00:00:00, UNBOUNDED)",
+            [
+                RangeQueryParameter(
+                    name="range_datetime",
+                    range_element_type="DATETIME",
+                    start=datetime.datetime(2016, 12, 5),
+                )
+            ],
+        ),
+        (
+            "SELECT @range_unbounded",
+            "[UNBOUNDED, UNBOUNDED)",
+            [
+                RangeQueryParameter(
+                    name="range_unbounded",
+                    range_element_type="DATETIME",
+                )
+            ],
+        ),
     ),
 )
 def test_query_parameters(