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

google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py

@@ -0,0 +1,239 @@
+# 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.
+
+
+# This file is automatically generated by CrossSync. Do not edit manually.
+
+from __future__ import annotations
+from typing import Any, Dict, Optional, Sequence, Tuple, TYPE_CHECKING
+from google.api_core import retry as retries
+from google.cloud.bigtable.data.execute_query._byte_cursor import _ByteCursor
+from google.cloud.bigtable.data._helpers import (
+    _attempt_timeout_generator,
+    _retry_exception_factory,
+)
+from google.cloud.bigtable.data.exceptions import (
+    EarlyMetadataCallError,
+    InvalidExecuteQueryResponse,
+)
+from google.cloud.bigtable.data.execute_query.values import QueryResultRow
+from google.cloud.bigtable.data.execute_query.metadata import Metadata
+from google.cloud.bigtable.data.execute_query._reader import (
+    _QueryResultRowReader,
+    _Reader,
+)
+from google.cloud.bigtable_v2.types.bigtable import (
+    ExecuteQueryRequest as ExecuteQueryRequestPB,
+    ExecuteQueryResponse,
+)
+from google.cloud.bigtable.data._cross_sync import CrossSync
+
+if TYPE_CHECKING:
+    from google.cloud.bigtable.data import BigtableDataClient as DataClientType
+
+
+def _has_resume_token(response: ExecuteQueryResponse) -> bool:
+    response_pb = response._pb
+    if response_pb.HasField("results"):
+        results = response_pb.results
+        return len(results.resume_token) > 0
+    return False
+
+
+class ExecuteQueryIterator:
+    def __init__(
+        self,
+        client: DataClientType,
+        instance_id: str,
+        app_profile_id: Optional[str],
+        request_body: Dict[str, Any],
+        prepare_metadata: Metadata,
+        attempt_timeout: float | None,
+        operation_timeout: float,
+        req_metadata: Sequence[Tuple[str, str]] = (),
+        retryable_excs: Sequence[type[Exception]] = (),
+    ) -> None:
+        """Collects responses from ExecuteQuery requests and parses them into QueryResultRows.
+
+        **Please Note** this is not meant to be constructed directly by applications. It should always
+        be created via the client. The constructor is subject to change.
+
+        It is **not thread-safe**. It should not be used by multiple threads.
+
+        Args:
+            client: bigtable client
+            instance_id: id of the instance on which the query is executed
+            request_body: dict representing the body of the ExecuteQueryRequest
+            attempt_timeout: the time budget for an individual network request, in seconds.
+                If it takes longer than this time to complete, the request will be cancelled with
+                a DeadlineExceeded exception, and a retry will be attempted.
+            operation_timeout: the time budget for the entire operation, in seconds.
+                Failed requests will be retried within the budget
+            req_metadata: metadata used while sending the gRPC request
+            retryable_excs: a list of errors that will be retried if encountered.
+        Raises:
+            None
+            :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
+        """
+        self._table_name = None
+        self._app_profile_id = app_profile_id
+        self._client = client
+        self._instance_id = instance_id
+        self._prepare_metadata = prepare_metadata
+        self._final_metadata = None
+        self._byte_cursor = _ByteCursor()
+        self._reader: _Reader[QueryResultRow] = _QueryResultRowReader()
+        self.has_received_token = False
+        self._result_generator = self._next_impl()
+        self._register_instance_task = None
+        self._is_closed = False
+        self._request_body = request_body
+        self._attempt_timeout_gen = _attempt_timeout_generator(
+            attempt_timeout, operation_timeout
+        )
+        self._stream = CrossSync._Sync_Impl.retry_target_stream(
+            self._make_request_with_resume_token,
+            retries.if_exception_type(*retryable_excs),
+            retries.exponential_sleep_generator(0.01, 60, multiplier=2),
+            operation_timeout,
+            exception_factory=_retry_exception_factory,
+        )
+        self._req_metadata = req_metadata
+        try:
+            self._register_instance_task = CrossSync._Sync_Impl.create_task(
+                self._client._register_instance,
+                self._instance_id,
+                self,
+                sync_executor=self._client._executor,
+            )
+        except RuntimeError as e:
+            raise RuntimeError(
+                f"{self.__class__.__name__} must be created within an async event loop context."
+            ) from e
+
+    @property
+    def is_closed(self) -> bool:
+        """Returns True if the iterator is closed, False otherwise."""
+        return self._is_closed
+
+    @property
+    def app_profile_id(self) -> Optional[str]:
+        """Returns the app_profile_id of the iterator."""
+        return self._app_profile_id
+
+    @property
+    def table_name(self) -> Optional[str]:
+        """Returns the table_name of the iterator."""
+        return self._table_name
+
+    def _make_request_with_resume_token(self):
+        """perfoms the rpc call using the correct resume token."""
+        resume_token = self._byte_cursor.prepare_for_new_request()
+        request = ExecuteQueryRequestPB(
+            {**self._request_body, "resume_token": resume_token}
+        )
+        return self._client._gapic_client.execute_query(
+            request,
+            timeout=next(self._attempt_timeout_gen),
+            metadata=self._req_metadata,
+            retry=None,
+        )
+
+    def _next_impl(self) -> CrossSync._Sync_Impl.Iterator[QueryResultRow]:
+        """Generator wrapping the response stream which parses the stream results
+        and returns full `QueryResultRow`s."""
+        for response in self._stream:
+            try:
+                if self._final_metadata is None and _has_resume_token(response):
+                    self._finalize_metadata()
+                batches_to_parse = self._byte_cursor.consume(response)
+                if not batches_to_parse:
+                    continue
+                if not self.metadata:
+                    raise ValueError(
+                        "Error parsing response before finalizing metadata"
+                    )
+                results = self._reader.consume(batches_to_parse, self.metadata)
+                if results is None:
+                    continue
+            except ValueError as e:
+                raise InvalidExecuteQueryResponse(
+                    "Invalid ExecuteQuery response received"
+                ) from e
+            for result in results:
+                yield result
+        if self._final_metadata is None:
+            self._finalize_metadata()
+        self.close()
+
+    def __next__(self) -> QueryResultRow:
+        """Yields QueryResultRows representing the results of the query.
+
+        :raises: :class:`ValueError <exceptions.ValueError>` as a safeguard if data is processed in an unexpected state
+        """
+        if self._is_closed:
+            raise CrossSync._Sync_Impl.StopIteration
+        return self._result_generator.__next__()
+
+    def __iter__(self):
+        return self
+
+    def _finalize_metadata(self) -> None:
+        """Sets _final_metadata to the metadata of the latest prepare_response.
+        The iterator should call this after either the first resume token is received or the
+        stream completes succesfully with no responses.
+
+        This can't be set on init because the metadata will be able to change due to plan refresh.
+        Plan refresh isn't implemented yet, but we want functionality to stay the same when it is.
+
+        For example the following scenario for query "SELECT * FROM table":
+          - Make a request, table has one column family 'cf'
+          - Return an incomplete batch
+          - request fails with transient error
+          - Meanwhile the table has had a second column family added 'cf2'
+          - Retry the request, get an error indicating the `prepared_query` has expired
+          - Refresh the prepared_query and retry the request, the new prepared_query
+            contains both 'cf' & 'cf2'
+          - It sends a new incomplete batch and resets the old outdated batch
+          - It send the next chunk with a checksum and resume_token, closing the batch.
+        In this we need to use the updated schema from the refreshed prepare request."""
+        self._final_metadata = self._prepare_metadata
+
+    @property
+    def metadata(self) -> Metadata:
+        """Returns query metadata from the server or None if the iterator has been closed
+        or if metadata has not been set yet.
+
+        Metadata will not be set until the first row has been yielded or response with no rows
+        completes.
+
+        raises: :class:`EarlyMetadataCallError` when called before the first row has been returned
+        or the iterator has completed with no rows in the response."""
+        if not self._final_metadata:
+            raise EarlyMetadataCallError()
+        return self._final_metadata
+
+    def close(self) -> None:
+        """Cancel all background tasks. Should be called all rows were processed.
+
+        :raises: :class:`ValueError <exceptions.ValueError>` if called in an invalid state
+        """
+        if self._is_closed:
+            return
+        if not self._byte_cursor.empty():
+            raise ValueError("Unexpected buffered data at end of executeQuery reqest")
+        self._is_closed = True
+        if self._register_instance_task is not None:
+            self._register_instance_task.cancel()
+        self._client._remove_instance_registration(self._instance_id, self)

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

@@ -0,0 +1,399 @@
+# 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.
+
+"""
+This module provides the SqlType class used for specifying types in
+ExecuteQuery and some utilities.
+
+The SqlTypes are used in Metadata returned by the ExecuteQuery operation as well
+as for specifying query parameter types explicitly.
+"""
+
+from collections import defaultdict
+import datetime
+from typing import Any, Dict, List, Optional, Set, Tuple, Type, Union
+
+from google.api_core.datetime_helpers import DatetimeWithNanoseconds
+from google.protobuf import timestamp_pb2  # type: ignore
+from google.type import date_pb2  # type: ignore
+
+from google.cloud.bigtable.data.execute_query.values import _NamedList
+from google.cloud.bigtable_v2 import ResultSetMetadata
+from google.cloud.bigtable_v2 import Type as PBType
+
+
+class SqlType:
+    """
+    Classes denoting types of values returned by Bigtable's ExecuteQuery operation.
+
+    Used in :class:`.Metadata`.
+    """
+
+    class Type:
+        expected_type: Optional[type] = None
+        value_pb_dict_field_name: Optional[str] = None
+        type_field_name: Optional[str] = None
+
+        @classmethod
+        def from_pb_type(cls, pb_type: Optional[PBType] = None):
+            return cls()
+
+        def _to_type_pb_dict(self) -> Dict[str, Any]:
+            if not self.type_field_name:
+                raise NotImplementedError(
+                    "Fill in expected_type and value_pb_dict_field_name"
+                )
+
+            return {self.type_field_name: {}}
+
+        def _to_value_pb_dict(self, value: Any) -> Dict[str, Any]:
+            if self.expected_type is None or self.value_pb_dict_field_name is None:
+                raise NotImplementedError(
+                    "Fill in expected_type and value_pb_dict_field_name"
+                )
+
+            if value is None:
+                return {}
+
+            if not isinstance(value, self.expected_type):
+                raise ValueError(
+                    f"Expected query parameter of type {self.expected_type.__name__}, got {type(value).__name__}"
+                )
+
+            return {self.value_pb_dict_field_name: value}
+
+        def __eq__(self, other):
+            return isinstance(other, type(self))
+
+        def __str__(self) -> str:
+            return self.__class__.__name__
+
+        def __repr__(self) -> str:
+            return self.__str__()
+
+    class Struct(_NamedList[Type], Type):
+        """Struct SQL type."""
+
+        @classmethod
+        def from_pb_type(cls, type_pb: Optional[PBType] = None) -> "SqlType.Struct":
+            if type_pb is None:
+                raise ValueError("missing required argument type_pb")
+            fields: List[Tuple[Optional[str], SqlType.Type]] = []
+            for field in type_pb.struct_type.fields:
+                fields.append((field.field_name, _pb_type_to_metadata_type(field.type)))
+            return cls(fields)
+
+        def _to_value_pb_dict(self, value: Any):
+            raise NotImplementedError("Struct is not supported as a query parameter")
+
+        def _to_type_pb_dict(self) -> Dict[str, Any]:
+            raise NotImplementedError("Struct is not supported as a query parameter")
+
+        def __eq__(self, other: object):
+            # Cannot use super() here - we'd either have to:
+            # - call super() in these base classes, which would in turn call Object.__eq__
+            #   to compare objects by identity and return a False, or
+            # - do not call super() in these base classes, which would result in calling only
+            #   one of the __eq__ methods (a super() in the base class would be required to call the other one), or
+            # - call super() in only one of the base classes, but that would be error prone and changing
+            #   the order of base classes would introduce unexpected behaviour.
+            # we also have to disable mypy because it doesn't see that SqlType.Struct == _NamedList[Type]
+            return SqlType.Type.__eq__(self, other) and _NamedList.__eq__(self, other)  # type: ignore
+
+        def __str__(self):
+            return super(_NamedList, self).__str__()
+
+    class Array(Type):
+        """Array SQL type."""
+
+        def __init__(self, element_type: "SqlType.Type"):
+            if isinstance(element_type, SqlType.Array):
+                raise ValueError("Arrays of arrays are not supported.")
+            if isinstance(element_type, SqlType.Map):
+                raise ValueError("Arrays of Maps are not supported.")
+            self._element_type = element_type
+
+        @property
+        def element_type(self):
+            return self._element_type
+
+        @classmethod
+        def from_pb_type(cls, type_pb: Optional[PBType] = None) -> "SqlType.Array":
+            if type_pb is None:
+                raise ValueError("missing required argument type_pb")
+            return cls(_pb_type_to_metadata_type(type_pb.array_type.element_type))
+
+        def _to_value_pb_dict(self, value: Any):
+            if value is None:
+                return {}
+
+            return {
+                "array_value": {
+                    "values": [
+                        self.element_type._to_value_pb_dict(entry) for entry in value
+                    ]
+                }
+            }
+
+        def _to_type_pb_dict(self) -> Dict[str, Any]:
+            return {
+                "array_type": {"element_type": self.element_type._to_type_pb_dict()}
+            }
+
+        def __eq__(self, other):
+            return super().__eq__(other) and self.element_type == other.element_type
+
+        def __str__(self) -> str:
+            return f"{self.__class__.__name__}<{str(self.element_type)}>"
+
+    class Map(Type):
+        """Map SQL type."""
+
+        def __init__(self, key_type: "SqlType.Type", value_type: "SqlType.Type"):
+            self._key_type = key_type
+            self._value_type = value_type
+
+        @property
+        def key_type(self):
+            return self._key_type
+
+        @property
+        def value_type(self):
+            return self._value_type
+
+        @classmethod
+        def from_pb_type(cls, type_pb: Optional[PBType] = None) -> "SqlType.Map":
+            if type_pb is None:
+                raise ValueError("missing required argument type_pb")
+            return cls(
+                _pb_type_to_metadata_type(type_pb.map_type.key_type),
+                _pb_type_to_metadata_type(type_pb.map_type.value_type),
+            )
+
+        def _to_type_pb_dict(self) -> Dict[str, Any]:
+            raise NotImplementedError("Map is not supported as a query parameter")
+
+        def _to_value_pb_dict(self, value: Any):
+            raise NotImplementedError("Map is not supported as a query parameter")
+
+        def __eq__(self, other):
+            return (
+                super().__eq__(other)
+                and self.key_type == other.key_type
+                and self.value_type == other.value_type
+            )
+
+        def __str__(self) -> str:
+            return (
+                f"{self.__class__.__name__}<"
+                f"{str(self._key_type)},{str(self._value_type)}>"
+            )
+
+    class Bytes(Type):
+        """Bytes SQL type."""
+
+        expected_type = bytes
+        value_pb_dict_field_name = "bytes_value"
+        type_field_name = "bytes_type"
+
+    class String(Type):
+        """String SQL type."""
+
+        expected_type = str
+        value_pb_dict_field_name = "string_value"
+        type_field_name = "string_type"
+
+    class Int64(Type):
+        """Int64 SQL type."""
+
+        expected_type = int
+        value_pb_dict_field_name = "int_value"
+        type_field_name = "int64_type"
+
+    class Float64(Type):
+        """Float64 SQL type."""
+
+        expected_type = float
+        value_pb_dict_field_name = "float_value"
+        type_field_name = "float64_type"
+
+    class Float32(Type):
+        """Float32 SQL type."""
+
+        expected_type = float
+        value_pb_dict_field_name = "float_value"
+        type_field_name = "float32_type"
+
+    class Bool(Type):
+        """Bool SQL type."""
+
+        expected_type = bool
+        value_pb_dict_field_name = "bool_value"
+        type_field_name = "bool_type"
+
+    class Timestamp(Type):
+        """
+        Timestamp SQL type.
+
+        Timestamp supports :class:`DatetimeWithNanoseconds` but Bigtable SQL does
+        not currently support nanoseconds precision. We support this for potential
+        compatibility in the future. Nanoseconds are currently ignored.
+        """
+
+        type_field_name = "timestamp_type"
+        expected_types = (
+            datetime.datetime,
+            DatetimeWithNanoseconds,
+        )
+
+        def _to_value_pb_dict(self, value: Any) -> Dict[str, Any]:
+            if value is None:
+                return {}
+
+            if not isinstance(value, self.expected_types):
+                raise ValueError(
+                    f"Expected one of {', '.join((_type.__name__ for _type in self.expected_types))}"
+                )
+
+            if isinstance(value, DatetimeWithNanoseconds):
+                return {"timestamp_value": value.timestamp_pb()}
+            else:  # value must be an instance of datetime.datetime
+                ts = timestamp_pb2.Timestamp()
+                ts.FromDatetime(value)
+                return {"timestamp_value": ts}
+
+    class Date(Type):
+        """Date SQL type."""
+
+        type_field_name = "date_type"
+        expected_type = datetime.date
+
+        def _to_value_pb_dict(self, value: Any) -> Dict[str, Any]:
+            if value is None:
+                return {}
+
+            if not isinstance(value, self.expected_type):
+                raise ValueError(
+                    f"Expected query parameter of type {self.expected_type.__name__}, got {type(value).__name__}"
+                )
+
+            return {
+                "date_value": date_pb2.Date(
+                    year=value.year,
+                    month=value.month,
+                    day=value.day,
+                )
+            }
+
+
+class Metadata:
+    """
+    Metadata class for the ExecuteQuery operation.
+
+    Args:
+        columns (List[Tuple[Optional[str], SqlType.Type]]): List of column
+            metadata tuples. Each tuple contains the column name and the column
+            type.
+    """
+
+    class Column:
+        def __init__(self, column_name: Optional[str], column_type: SqlType.Type):
+            self._column_name = column_name
+            self._column_type = column_type
+
+        @property
+        def column_name(self) -> Optional[str]:
+            return self._column_name
+
+        @property
+        def column_type(self) -> SqlType.Type:
+            return self._column_type
+
+    @property
+    def columns(self) -> List[Column]:
+        return self._columns
+
+    def __init__(
+        self, columns: Optional[List[Tuple[Optional[str], SqlType.Type]]] = None
+    ):
+        self._columns: List[Metadata.Column] = []
+        self._column_indexes: Dict[str, List[int]] = defaultdict(list)
+        self._duplicate_names: Set[str] = set()
+
+        if columns:
+            for column_name, column_type in columns:
+                if column_name is not None:
+                    if column_name in self._column_indexes:
+                        self._duplicate_names.add(column_name)
+                    self._column_indexes[column_name].append(len(self._columns))
+                self._columns.append(Metadata.Column(column_name, column_type))
+
+    def __getitem__(self, index_or_name: Union[str, int]) -> Column:
+        if isinstance(index_or_name, str):
+            if index_or_name in self._duplicate_names:
+                raise KeyError(
+                    f"Ambigious column name: '{index_or_name}', use index instead."
+                    f" Field present on indexes {', '.join(map(str, self._column_indexes[index_or_name]))}."
+                )
+            if index_or_name not in self._column_indexes:
+                raise KeyError(f"No such column: {index_or_name}")
+            index = self._column_indexes[index_or_name][0]
+        else:
+            index = index_or_name
+        return self._columns[index]
+
+    def __len__(self):
+        return len(self._columns)
+
+    def __str__(self) -> str:
+        columns_str = ", ".join([str(column) for column in self._columns])
+        return f"{self.__class__.__name__}([{columns_str}])"
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+def _pb_metadata_to_metadata_types(
+    metadata_pb: ResultSetMetadata,
+) -> Metadata:
+    if "proto_schema" in metadata_pb:
+        fields: List[Tuple[Optional[str], SqlType.Type]] = []
+        for column_metadata in metadata_pb.proto_schema.columns:
+            fields.append(
+                (column_metadata.name, _pb_type_to_metadata_type(column_metadata.type))
+            )
+        return Metadata(fields)
+    raise ValueError("Invalid ResultSetMetadata object received.")
+
+
+_PROTO_TYPE_TO_METADATA_TYPE_FACTORY: Dict[str, Type[SqlType.Type]] = {
+    "bytes_type": SqlType.Bytes,
+    "string_type": SqlType.String,
+    "int64_type": SqlType.Int64,
+    "float32_type": SqlType.Float32,
+    "float64_type": SqlType.Float64,
+    "bool_type": SqlType.Bool,
+    "timestamp_type": SqlType.Timestamp,
+    "date_type": SqlType.Date,
+    "struct_type": SqlType.Struct,
+    "array_type": SqlType.Array,
+    "map_type": SqlType.Map,
+}
+
+
+def _pb_type_to_metadata_type(type_pb: PBType) -> SqlType.Type:
+    kind = PBType.pb(type_pb).WhichOneof("kind")
+    if kind in _PROTO_TYPE_TO_METADATA_TYPE_FACTORY:
+        return _PROTO_TYPE_TO_METADATA_TYPE_FACTORY[kind].from_pb_type(type_pb)
+    raise ValueError(f"Unrecognized response data type: {type_pb}")