google-cloud-bigtable 2.30.0__py3-none-any.whl

google/cloud/bigtable/data/execute_query/_byte_cursor.py

@@ -0,0 +1,123 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional
+
+from google.cloud.bigtable.data.execute_query._checksum import _CRC32C
+from google.cloud.bigtable_v2 import ExecuteQueryResponse
+
+
+class _ByteCursor:
+    """
+    Buffers bytes from `ExecuteQuery` responses until resume_token is received or end-of-stream
+    is reached. :class:`google.cloud.bigtable_v2.types.bigtable.ExecuteQueryResponse` obtained from
+    the server should be passed to the ``consume`` method and its non-None results should be passed
+    to appropriate :class:`google.cloud.bigtable.execute_query_reader._Reader` for parsing gathered
+    bytes.
+
+    This class consumes data obtained externally to be usable in both sync and async clients.
+
+    See :class:`google.cloud.bigtable.execute_query_reader._Reader` for more context.
+    """
+
+    def __init__(self):
+        self._batch_buffer = bytearray()
+        self._batches: List[bytes] = []
+        self._resume_token = None
+
+    def reset(self):
+        self._batch_buffer = bytearray()
+        self._batches = []
+
+    def prepare_for_new_request(self):
+        """
+        Prepares this ``_ByteCursor`` for retrying an ``ExecuteQuery`` request.
+
+        Clears internal buffers of this ``_ByteCursor`` and returns last received
+        ``resume_token`` to be used in retried request.
+
+        This is the only method that returns ``resume_token`` to the user.
+        Returning the token to the user is tightly coupled with clearing internal
+        buffers to prevent accidental retry without clearing the state, what would
+        cause invalid results. ``resume_token`` are not needed in other cases,
+        thus they is no separate getter for it.
+
+        Returns:
+            bytes: Last received resume_token.
+        """
+        # The first response of any retried stream will always contain reset, so
+        # this isn't actually necessary, but we do it for safety
+        self.reset()
+        return self._resume_token
+
+    def empty(self) -> bool:
+        return not self._batch_buffer and not self._batches
+
+    def consume(self, response: ExecuteQueryResponse) -> Optional[List[bytes]]:
+        """
+        Reads results bytes from an ``ExecuteQuery`` response and adds them to a buffer.
+
+        If the response contains a ``resume_token``:
+        - the ``resume_token`` is saved in this ``_ByteCursor``, and
+        - internal buffers are flushed and returned to the caller.
+
+        ``resume_token`` is not available directly, but can be retrieved by calling
+        :meth:`._ByteCursor.prepare_for_new_request` when preparing to retry a request.
+
+        Args:
+            response (google.cloud.bigtable_v2.types.bigtable.ExecuteQueryResponse):
+                Response obtained from the stream.
+
+        Returns:
+            bytes or None: List of bytes if buffers were flushed or None otherwise.
+            Each element in the list represents the bytes of a `ProtoRows` message.
+
+        Raises:
+            ValueError: If provided ``ExecuteQueryResponse`` is not valid
+                or contains bytes representing response of a different kind than previously
+                processed responses.
+        """
+        response_pb = response._pb  # proto-plus attribute retrieval is slow.
+
+        if response_pb.HasField("results"):
+            results = response_pb.results
+            if results.reset:
+                self.reset()
+            if results.HasField("proto_rows_batch"):
+                self._batch_buffer.extend(results.proto_rows_batch.batch_data)
+                # Note that 0 is a valid checksum so we must check for field presence
+                if results.HasField("batch_checksum"):
+                    expected_checksum = results.batch_checksum
+                    checksum = _CRC32C.checksum(self._batch_buffer)
+                    if expected_checksum != checksum:
+                        raise ValueError(
+                            f"Unexpected checksum mismatch. Expected: {expected_checksum}, got: {checksum}"
+                        )
+                    # We have a complete batch so we move it to batches and reset the
+                    # batch_buffer
+                    self._batches.append(memoryview(self._batch_buffer))
+                    self._batch_buffer = bytearray()
+
+            if results.resume_token:
+                self._resume_token = results.resume_token
+
+                if self._batches:
+                    if self._batch_buffer:
+                        raise ValueError("Unexpected resume_token without checksum")
+                    return_value = self._batches
+                    self._batches = []
+                    return return_value
+        else:
+            raise ValueError(f"Unexpected ExecuteQueryResponse: {response}")
+        return None

google/cloud/bigtable/data/execute_query/_checksum.py

@@ -0,0 +1,43 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import warnings
+
+with warnings.catch_warnings(record=True) as import_warning:
+    import google_crc32c  # type: ignore
+
+
+class _CRC32C(object):
+    """
+    Wrapper around ``google_crc32c`` library
+    """
+
+    warn_emitted = False
+
+    @classmethod
+    def checksum(cls, val: bytearray) -> int:
+        """
+        Returns the crc32c checksum of the data.
+        """
+        if import_warning and not cls.warn_emitted:
+            cls.warn_emitted = True
+            warnings.warn(
+                "Using pure python implementation of `google-crc32` for ExecuteQuery response "
+                "validation. This is significantly slower than the c extension. If possible, "
+                "run in an environment that supports the c extension.",
+                RuntimeWarning,
+            )
+        memory_view = memoryview(val)
+        return google_crc32c.value(bytes(memory_view))

google/cloud/bigtable/data/execute_query/_parameters_formatting.py

@@ -0,0 +1,155 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+from typing import Any, Dict, Optional
+
+from google.api_core.datetime_helpers import DatetimeWithNanoseconds
+
+from google.cloud.bigtable.data.exceptions import ParameterTypeInferenceFailed
+from google.cloud.bigtable.data.execute_query.metadata import SqlType
+from google.cloud.bigtable.data.execute_query.values import ExecuteQueryValueType
+from google.cloud.bigtable_v2.types.data import Value
+
+
+def _format_execute_query_params(
+    params: Optional[Dict[str, ExecuteQueryValueType]],
+    parameter_types: Optional[Dict[str, SqlType.Type]],
+) -> Dict[str, Value]:
+    """
+    Takes a dictionary of param_name -> param_value and optionally parameter types.
+    If the parameters types are not provided, this function tries to infer them.
+
+    Args:
+        params (Optional[Dict[str, ExecuteQueryValueType]]): mapping from parameter names
+        like they appear in query (without @ at the beginning) to their values.
+        Only values of type ExecuteQueryValueType are permitted.
+        parameter_types (Optional[Dict[str, SqlType.Type]]): mapping of parameter names
+        to their types.
+
+    Raises:
+        ValueError: raised when parameter types cannot be inferred and were not
+        provided explicitly.
+
+    Returns:
+         dictionary prasable to a protobuf represenging parameters as defined
+         in ExecuteQueryRequest.params
+    """
+    if not params:
+        return {}
+    parameter_types = parameter_types or {}
+
+    result_values = {}
+    for key, value in params.items():
+        user_provided_type = parameter_types.get(key)
+        try:
+            if user_provided_type:
+                if not isinstance(user_provided_type, SqlType.Type):
+                    raise ValueError(
+                        f"Parameter type for {key} should be provided as an instance of SqlType.Type subclass."
+                    )
+                param_type = user_provided_type
+            else:
+                param_type = _detect_type(value)
+
+            value_pb_dict = _convert_value_to_pb_value_dict(value, param_type)
+        except ValueError as err:
+            raise ValueError(f"Error when parsing parameter {key}") from err
+        result_values[key] = value_pb_dict
+
+    return result_values
+
+
+def _to_param_types(
+    params: Optional[Dict[str, ExecuteQueryValueType]],
+    param_types: Optional[Dict[str, SqlType.Type]],
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Takes the params and user supplied types and creates a param_type dict for the PrepareQuery api
+
+    Args:
+        params: Dict of param name to param value
+        param_types: Dict of param name to param type for params with types that cannot be inferred
+
+    Returns:
+        Dict containing the param name and type for each parameter
+    """
+    if params is None:
+        return {}
+    formatted_types = {}
+    for param_key, param_value in params.items():
+        if param_types and param_key in param_types:
+            formatted_types[param_key] = param_types[param_key]._to_type_pb_dict()
+        else:
+            formatted_types[param_key] = _detect_type(param_value)._to_type_pb_dict()
+    return formatted_types
+
+
+def _convert_value_to_pb_value_dict(
+    value: ExecuteQueryValueType, param_type: SqlType.Type
+) -> Any:
+    """
+    Takes a value and converts it to a dictionary parsable to a protobuf.
+
+    Args:
+        value (ExecuteQueryValueType): value
+        param_type (SqlType.Type): object describing which ExecuteQuery type the value represents.
+
+    Returns:
+        dictionary parsable to a protobuf.
+    """
+    # type field will be set only in top-level Value.
+    value_dict = param_type._to_value_pb_dict(value)
+    value_dict["type_"] = param_type._to_type_pb_dict()
+    return value_dict
+
+
+_TYPES_TO_TYPE_DICTS = [
+    (bytes, SqlType.Bytes()),
+    (str, SqlType.String()),
+    (bool, SqlType.Bool()),
+    (int, SqlType.Int64()),
+    (DatetimeWithNanoseconds, SqlType.Timestamp()),
+    (datetime.datetime, SqlType.Timestamp()),
+    (datetime.date, SqlType.Date()),
+]
+
+
+def _detect_type(value: ExecuteQueryValueType) -> SqlType.Type:
+    """
+    Infers the ExecuteQuery type based on value. Raises error if type is amiguous.
+    raises ParameterTypeInferenceFailed if not possible.
+    """
+    if value is None:
+        raise ParameterTypeInferenceFailed(
+            "Cannot infer type of None, please provide the type manually."
+        )
+
+    if isinstance(value, list):
+        raise ParameterTypeInferenceFailed(
+            "Cannot infer type of ARRAY parameters, please provide the type manually."
+        )
+
+    if isinstance(value, float):
+        raise ParameterTypeInferenceFailed(
+            "Cannot infer type of float, must specify either FLOAT32 or FLOAT64 type manually."
+        )
+
+    for field_type, type_dict in _TYPES_TO_TYPE_DICTS:
+        if isinstance(value, field_type):
+            return type_dict
+
+    raise ParameterTypeInferenceFailed(
+        f"Cannot infer type of {type(value).__name__}, please provide the type manually."
+    )

google/cloud/bigtable/data/execute_query/_query_result_parsing_utils.py

@@ -0,0 +1,134 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Callable, Dict, Type
+from google.cloud.bigtable.data.execute_query.values import Struct
+from google.cloud.bigtable.data.execute_query.metadata import SqlType
+from google.cloud.bigtable_v2 import Value as PBValue
+from google.api_core.datetime_helpers import DatetimeWithNanoseconds
+
+_REQUIRED_PROTO_FIELDS = {
+    SqlType.Bytes: "bytes_value",
+    SqlType.String: "string_value",
+    SqlType.Int64: "int_value",
+    SqlType.Float32: "float_value",
+    SqlType.Float64: "float_value",
+    SqlType.Bool: "bool_value",
+    SqlType.Timestamp: "timestamp_value",
+    SqlType.Date: "date_value",
+    SqlType.Struct: "array_value",
+    SqlType.Array: "array_value",
+    SqlType.Map: "array_value",
+}
+
+
+def _parse_array_type(value: PBValue, metadata_type: SqlType.Array) -> Any:
+    """
+    used for parsing an array represented as a protobuf to a python list.
+    """
+    return list(
+        map(
+            lambda val: _parse_pb_value_to_python_value(
+                val, metadata_type.element_type
+            ),
+            value.array_value.values,
+        )
+    )
+
+
+def _parse_map_type(value: PBValue, metadata_type: SqlType.Map) -> Any:
+    """
+    used for parsing a map represented as a protobuf to a python dict.
+
+    Values of type `Map` are stored in a `Value.array_value` where each entry
+    is another `Value.array_value` with two elements (the key and the value,
+    in that order).
+    Normally encoded Map values won't have repeated keys, however, the client
+    must handle the case in which they do. If the same key appears
+    multiple times, the _last_ value takes precedence.
+    """
+
+    try:
+        return dict(
+            map(
+                lambda map_entry: (
+                    _parse_pb_value_to_python_value(
+                        map_entry.array_value.values[0], metadata_type.key_type
+                    ),
+                    _parse_pb_value_to_python_value(
+                        map_entry.array_value.values[1], metadata_type.value_type
+                    ),
+                ),
+                value.array_value.values,
+            )
+        )
+    except IndexError:
+        raise ValueError("Invalid map entry - less or more than two values.")
+
+
+def _parse_struct_type(value: PBValue, metadata_type: SqlType.Struct) -> Struct:
+    """
+    used for parsing a struct represented as a protobuf to a
+    google.cloud.bigtable.data.execute_query.Struct
+    """
+    if len(value.array_value.values) != len(metadata_type.fields):
+        raise ValueError("Mismatched lengths of values and types.")
+
+    struct = Struct()
+    for value, field in zip(value.array_value.values, metadata_type.fields):
+        field_name, field_type = field
+        struct.add_field(field_name, _parse_pb_value_to_python_value(value, field_type))
+
+    return struct
+
+
+def _parse_timestamp_type(
+    value: PBValue, metadata_type: SqlType.Timestamp
+) -> DatetimeWithNanoseconds:
+    """
+    used for parsing a timestamp represented as a protobuf to DatetimeWithNanoseconds
+    """
+    return DatetimeWithNanoseconds.from_timestamp_pb(value.timestamp_value)
+
+
+_TYPE_PARSERS: Dict[Type[SqlType.Type], Callable[[PBValue, Any], Any]] = {
+    SqlType.Timestamp: _parse_timestamp_type,
+    SqlType.Struct: _parse_struct_type,
+    SqlType.Array: _parse_array_type,
+    SqlType.Map: _parse_map_type,
+}
+
+
+def _parse_pb_value_to_python_value(value: PBValue, metadata_type: SqlType.Type) -> Any:
+    """
+    used for converting the value represented as a protobufs to a python object.
+    """
+    value_kind = value.WhichOneof("kind")
+    if not value_kind:
+        return None
+
+    kind = type(metadata_type)
+    if not value.HasField(_REQUIRED_PROTO_FIELDS[kind]):
+        raise ValueError(
+            f"{_REQUIRED_PROTO_FIELDS[kind]} field for {kind.__name__} type not found in a Value."
+        )
+
+    if kind in _TYPE_PARSERS:
+        parser = _TYPE_PARSERS[kind]
+        return parser(value, metadata_type)
+    elif kind in _REQUIRED_PROTO_FIELDS:
+        field_name = _REQUIRED_PROTO_FIELDS[kind]
+        return getattr(value, field_name)
+    else:
+        raise ValueError(f"Unknown kind {kind}")

google/cloud/bigtable/data/execute_query/_reader.py

@@ -0,0 +1,122 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import (
+    List,
+    TypeVar,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+)
+from abc import ABC, abstractmethod
+
+from google.cloud.bigtable_v2 import ProtoRows, Value as PBValue
+
+from google.cloud.bigtable.data.execute_query._query_result_parsing_utils import (
+    _parse_pb_value_to_python_value,
+)
+
+from google.cloud.bigtable.helpers import batched
+
+from google.cloud.bigtable.data.execute_query.values import QueryResultRow
+from google.cloud.bigtable.data.execute_query.metadata import Metadata
+
+
+T = TypeVar("T")
+
+
+class _Reader(ABC, Generic[T]):
+    """
+    An interface for classes that consume and parse bytes returned by ``_ByteCursor``.
+    Parsed bytes should be gathered into bundles (rows or columns) of expected size
+    and converted to an appropriate type ``T`` that will be returned as a semantically
+    meaningful result to the library user by
+    :meth:`google.cloud.bigtable.instance.Instance.execute_query` or
+    :meth:`google.cloud.bigtable.data._async.client.BigtableDataClientAsync.execute_query`
+    methods.
+
+    This class consumes data obtained externally to be usable in both sync and async clients.
+
+    See :class:`google.cloud.bigtable.byte_cursor._ByteCursor` for more context.
+    """
+
+    @abstractmethod
+    def consume(
+        self, batches_to_consume: List[bytes], metadata: Metadata
+    ) -> Optional[Iterable[T]]:
+        """This method receives a list of batches of bytes to be parsed as ProtoRows messages.
+        It then uses the metadata to group the values in the parsed messages into rows. Returns
+        None if batches_to_consume is empty
+        Args:
+            bytes_to_consume (bytes): chunk of parsable byte batches received from
+                :meth:`google.cloud.bigtable.byte_cursor._ByteCursor.consume`
+                method.
+            metadata: metadata used to transform values to rows
+
+        Returns:
+            Iterable[T] or None: Iterable if gathered values can form one or more instances of T,
+                or None if there is not enough data to construct at least one instance of T with
+                appropriate number of entries.
+        """
+        raise NotImplementedError
+
+
+class _QueryResultRowReader(_Reader[QueryResultRow]):
+    """
+    A :class:`._Reader` consuming bytes representing
+    :class:`google.cloud.bigtable_v2.types.Type`
+    and producing :class:`google.cloud.bigtable.execute_query.QueryResultRow`.
+
+    Number of entries in each row is determined by number of columns in
+    :class:`google.cloud.bigtable.execute_query.Metadata` obtained from
+    :class:`google.cloud.bigtable.byte_cursor._ByteCursor` passed in the constructor.
+    """
+
+    def _parse_proto_rows(self, bytes_to_parse: bytes) -> Iterable[PBValue]:
+        proto_rows = ProtoRows.pb().FromString(bytes_to_parse)
+        return proto_rows.values
+
+    def _construct_query_result_row(
+        self, values: Sequence[PBValue], metadata: Metadata
+    ) -> QueryResultRow:
+        result = QueryResultRow()
+        columns = metadata.columns
+
+        assert len(values) == len(
+            columns
+        ), "This function should be called only when count of values matches count of columns."
+
+        for column, value in zip(columns, values):
+            parsed_value = _parse_pb_value_to_python_value(value, column.column_type)
+            result.add_field(column.column_name, parsed_value)
+        return result
+
+    def consume(
+        self, batches_to_consume: List[bytes], metadata: Metadata
+    ) -> Optional[Iterable[QueryResultRow]]:
+        num_columns = len(metadata.columns)
+        rows = []
+        for batch_bytes in batches_to_consume:
+            values = self._parse_proto_rows(batch_bytes)
+            for row_data in batched(values, n=num_columns):
+                if len(row_data) == num_columns:
+                    rows.append(self._construct_query_result_row(row_data, metadata))
+                else:
+                    raise ValueError(
+                        "Unexpected error, recieved bad number of values. "
+                        f"Expected {num_columns} got {len(row_data)}."
+                    )
+
+        return rows