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

google/cloud/bigtable/data/execute_query/values.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 collections import defaultdict
+from typing import (
+    Optional,
+    List,
+    Dict,
+    Set,
+    Union,
+    TypeVar,
+    Generic,
+    Tuple,
+    Mapping,
+)
+from google.type import date_pb2  # type: ignore
+from google.api_core.datetime_helpers import DatetimeWithNanoseconds
+
+T = TypeVar("T")
+
+
+class _NamedList(Generic[T]):
+    """
+    A class designed to store a list of elements, which can be accessed by
+    name or index.
+    This class is different from namedtuple, because namedtuple has some
+    restrictions on names of fields and we do not want to have them.
+    """
+
+    _str_cls_name = "_NamedList"
+
+    def __init__(self, fields: Optional[List[Tuple[Optional[str], T]]] = None):
+        self._fields: List[Tuple[Optional[str], T]] = []
+        self._field_indexes: Dict[str, List[int]] = defaultdict(list)
+        self._duplicate_names: Set[str] = set()
+
+        if fields:
+            for field_name, field_type in fields:
+                self.add_field(field_name, field_type)
+
+    def add_field(self, name: Optional[str], value: T):
+        if name:
+            if name in self._field_indexes:
+                self._duplicate_names.add(name)
+            self._field_indexes[name].append(len(self._fields))
+        self._fields.append((name, value))
+
+    @property
+    def fields(self):
+        return self._fields
+
+    def __getitem__(self, index_or_name: Union[str, int]):
+        if isinstance(index_or_name, str):
+            if index_or_name in self._duplicate_names:
+                raise KeyError(
+                    f"Ambigious field name: '{index_or_name}', use index instead."
+                    f" Field present on indexes {', '.join(map(str, self._field_indexes[index_or_name]))}."
+                )
+            if index_or_name not in self._field_indexes:
+                raise KeyError(f"No such field: {index_or_name}")
+            index = self._field_indexes[index_or_name][0]
+        else:
+            index = index_or_name
+        return self._fields[index][1]
+
+    def __len__(self):
+        return len(self._fields)
+
+    def __eq__(self, other):
+        if not isinstance(other, _NamedList):
+            return False
+
+        return (
+            self._fields == other._fields
+            and self._field_indexes == other._field_indexes
+        )
+
+    def __str__(self) -> str:
+        fields_str = ", ".join([str(field) for field in self._fields])
+        return f"{self.__class__.__name__}([{fields_str}])"
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+
+ExecuteQueryValueType = Union[
+    int,
+    float,
+    bool,
+    bytes,
+    str,
+    # Note that Bigtable SQL does not currently support nanosecond precision,
+    # only microseconds. We use this for compatibility with potential future
+    # support
+    DatetimeWithNanoseconds,
+    date_pb2.Date,
+    "Struct",
+    List["ExecuteQueryValueType"],
+    Mapping[Union[str, int, bytes], "ExecuteQueryValueType"],
+]
+
+
+class QueryResultRow(_NamedList[ExecuteQueryValueType]):
+    """
+    Represents a single row of the result
+    """
+
+
+class Struct(_NamedList[ExecuteQueryValueType]):
+    """
+    Represents a struct value in the result
+    """

google/cloud/bigtable/data/mutations.py

@@ -0,0 +1,380 @@
+# Copyright 2023 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 __future__ import annotations
+from typing import Any
+import time
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+from sys import getsizeof
+
+import google.cloud.bigtable_v2.types.bigtable as types_pb
+import google.cloud.bigtable_v2.types.data as data_pb
+
+from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE
+
+
+# special value for SetCell mutation timestamps. If set, server will assign a timestamp
+_SERVER_SIDE_TIMESTAMP = -1
+
+# mutation entries above this should be rejected
+_MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000
+
+
+class Mutation(ABC):
+    """
+    Abstract base class for mutations.
+
+    This class defines the interface for different types of mutations that can be
+    applied to Bigtable rows.
+    """
+
+    @abstractmethod
+    def _to_dict(self) -> dict[str, Any]:
+        """
+        Convert the mutation to a dictionary representation.
+
+        Returns:
+            dict[str, Any]: A dictionary representation of the mutation.
+        """
+        raise NotImplementedError
+
+    def _to_pb(self) -> data_pb.Mutation:
+        """
+        Convert the mutation to a protobuf representation.
+
+        Returns:
+            Mutation: A protobuf representation of the mutation.
+        """
+        return data_pb.Mutation(**self._to_dict())
+
+    def is_idempotent(self) -> bool:
+        """
+        Check if the mutation is idempotent
+
+        Idempotent mutations can be safely retried on failure.
+
+        Returns:
+            bool: True if the mutation is idempotent, False otherwise.
+        """
+        return True
+
+    def __str__(self) -> str:
+        """
+        Return a string representation of the mutation.
+
+        Returns:
+            str: A string representation of the mutation.
+        """
+        return str(self._to_dict())
+
+    def size(self) -> int:
+        """
+        Get the size of the mutation in bytes
+
+        Returns:
+            int: The size of the mutation in bytes.
+        """
+        return getsizeof(self._to_dict())
+
+    @classmethod
+    def _from_dict(cls, input_dict: dict[str, Any]) -> Mutation:
+        """
+        Create a `Mutation` instance from a dictionary representation.
+
+        Args:
+            input_dict: A dictionary representation of the mutation.
+        Returns:
+            Mutation: A Mutation instance created from the dictionary.
+        Raises:
+            ValueError: If the input dictionary is invalid or does not represent a valid mutation type.
+        """
+        instance: Mutation | None = None
+        try:
+            if "set_cell" in input_dict:
+                details = input_dict["set_cell"]
+                instance = SetCell(
+                    details["family_name"],
+                    details["column_qualifier"],
+                    details["value"],
+                    details["timestamp_micros"],
+                )
+            elif "delete_from_column" in input_dict:
+                details = input_dict["delete_from_column"]
+                time_range = details.get("time_range", {})
+                start = time_range.get("start_timestamp_micros", None)
+                end = time_range.get("end_timestamp_micros", None)
+                instance = DeleteRangeFromColumn(
+                    details["family_name"], details["column_qualifier"], start, end
+                )
+            elif "delete_from_family" in input_dict:
+                details = input_dict["delete_from_family"]
+                instance = DeleteAllFromFamily(details["family_name"])
+            elif "delete_from_row" in input_dict:
+                instance = DeleteAllFromRow()
+        except KeyError as e:
+            raise ValueError("Invalid mutation dictionary") from e
+        if instance is None:
+            raise ValueError("No valid mutation found")
+        if not issubclass(instance.__class__, cls):
+            raise ValueError("Mutation type mismatch")
+        return instance
+
+
+class SetCell(Mutation):
+    """
+    Mutation to set the value of a cell.
+
+    Args:
+        family: The name of the column family to which the new cell belongs.
+        qualifier: The column qualifier of the new cell.
+        new_value: The value of the new cell.
+        timestamp_micros: The timestamp of the new cell. If `None`,
+            the current timestamp will be used. Timestamps will be sent with
+            millisecond precision. Extra precision will be truncated. If -1, the
+            server will assign a timestamp. Note that `SetCell` mutations with
+            server-side timestamps are non-idempotent operations and will not be retried.
+
+    Raises:
+        TypeError: If `qualifier` is not `bytes` or `str`.
+        TypeError: If `new_value` is not `bytes`, `str`, or `int`.
+        ValueError: If `timestamp_micros` is less than `_SERVER_SIDE_TIMESTAMP`.
+    """
+
+    def __init__(
+        self,
+        family: str,
+        qualifier: bytes | str,
+        new_value: bytes | str | int,
+        timestamp_micros: int | None = None,
+    ):
+        qualifier = qualifier.encode() if isinstance(qualifier, str) else qualifier
+        if not isinstance(qualifier, bytes):
+            raise TypeError("qualifier must be bytes or str")
+        if isinstance(new_value, str):
+            new_value = new_value.encode()
+        elif isinstance(new_value, int):
+            if abs(new_value) > _MAX_INCREMENT_VALUE:
+                raise ValueError(
+                    "int values must be between -2**63 and 2**63 (64-bit signed int)"
+                )
+            new_value = new_value.to_bytes(8, "big", signed=True)
+        if not isinstance(new_value, bytes):
+            raise TypeError("new_value must be bytes, str, or int")
+        if timestamp_micros is None:
+            # use current timestamp, with milisecond precision
+            timestamp_micros = time.time_ns() // 1000
+            timestamp_micros = timestamp_micros - (timestamp_micros % 1000)
+        if timestamp_micros < _SERVER_SIDE_TIMESTAMP:
+            raise ValueError(
+                f"timestamp_micros must be positive (or {_SERVER_SIDE_TIMESTAMP} for server-side timestamp)"
+            )
+        self.family = family
+        self.qualifier = qualifier
+        self.new_value = new_value
+        self.timestamp_micros = timestamp_micros
+
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "set_cell": {
+                "family_name": self.family,
+                "column_qualifier": self.qualifier,
+                "timestamp_micros": self.timestamp_micros,
+                "value": self.new_value,
+            }
+        }
+
+    def is_idempotent(self) -> bool:
+        return self.timestamp_micros != _SERVER_SIDE_TIMESTAMP
+
+
+@dataclass
+class DeleteRangeFromColumn(Mutation):
+    """
+    Mutation to delete a range of cells from a column.
+
+    Args:
+        family: The name of the column family.
+         qualifier: The column qualifier.
+        start_timestamp_micros: The start timestamp of the range to
+            delete. `None` represents 0. Defaults to `None`.
+        end_timestamp_micros: The end timestamp of the range to
+            delete. `None` represents infinity. Defaults to `None`.
+    Raises:
+        ValueError: If `start_timestamp_micros` is greater than `end_timestamp_micros`.
+    """
+
+    family: str
+    qualifier: bytes
+    # None represents 0
+    start_timestamp_micros: int | None = None
+    # None represents infinity
+    end_timestamp_micros: int | None = None
+
+    def __post_init__(self):
+        if (
+            self.start_timestamp_micros is not None
+            and self.end_timestamp_micros is not None
+            and self.start_timestamp_micros > self.end_timestamp_micros
+        ):
+            raise ValueError("start_timestamp_micros must be <= end_timestamp_micros")
+
+    def _to_dict(self) -> dict[str, Any]:
+        timestamp_range = {}
+        if self.start_timestamp_micros is not None:
+            timestamp_range["start_timestamp_micros"] = self.start_timestamp_micros
+        if self.end_timestamp_micros is not None:
+            timestamp_range["end_timestamp_micros"] = self.end_timestamp_micros
+        return {
+            "delete_from_column": {
+                "family_name": self.family,
+                "column_qualifier": self.qualifier,
+                "time_range": timestamp_range,
+            }
+        }
+
+
+@dataclass
+class DeleteAllFromFamily(Mutation):
+    """
+    Mutation to delete all cells from a column family.
+
+    Args:
+        family_to_delete: The name of the column family to delete.
+    """
+
+    family_to_delete: str
+
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "delete_from_family": {
+                "family_name": self.family_to_delete,
+            }
+        }
+
+
+@dataclass
+class DeleteAllFromRow(Mutation):
+    """
+    Mutation to delete all cells from a row.
+    """
+
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "delete_from_row": {},
+        }
+
+
+class RowMutationEntry:
+    """
+    A single entry in a `MutateRows` request.
+
+    This class represents a set of mutations to apply to a specific row in a
+    Bigtable table.
+
+    Args:
+        row_key: The key of the row to mutate.
+        mutations: The mutation or list of mutations to apply
+            to the row.
+
+    Raises:
+        ValueError: If `mutations` is empty or contains more than
+            `_MUTATE_ROWS_REQUEST_MUTATION_LIMIT` mutations.
+    """
+
+    def __init__(self, row_key: bytes | str, mutations: Mutation | list[Mutation]):
+        if isinstance(row_key, str):
+            row_key = row_key.encode("utf-8")
+        if isinstance(mutations, Mutation):
+            mutations = [mutations]
+        if len(mutations) == 0:
+            raise ValueError("mutations must not be empty")
+        elif len(mutations) > _MUTATE_ROWS_REQUEST_MUTATION_LIMIT:
+            raise ValueError(
+                f"entries must have <= {_MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations"
+            )
+        self.row_key = row_key
+        self.mutations = tuple(mutations)
+
+    def _to_dict(self) -> dict[str, Any]:
+        """
+        Convert the mutation entry to a dictionary representation.
+
+        Returns:
+            dict[str, Any]: A dictionary representation of the mutation entry
+        """
+        return {
+            "row_key": self.row_key,
+            "mutations": [mutation._to_dict() for mutation in self.mutations],
+        }
+
+    def _to_pb(self) -> types_pb.MutateRowsRequest.Entry:
+        """
+        Convert the mutation entry to a protobuf representation.
+
+        Returns:
+            MutateRowsRequest.Entry: A protobuf representation of the mutation entry.
+        """
+        return types_pb.MutateRowsRequest.Entry(
+            row_key=self.row_key,
+            mutations=[mutation._to_pb() for mutation in self.mutations],
+        )
+
+    def is_idempotent(self) -> bool:
+        """
+        Check if all mutations in the entry are idempotent.
+
+        Returns:
+            bool: True if all mutations in the entry are idempotent, False otherwise.
+        """
+        return all(mutation.is_idempotent() for mutation in self.mutations)
+
+    def size(self) -> int:
+        """
+        Get the size of the mutation entry in bytes.
+
+        Returns:
+            int: The size of the mutation entry in bytes.
+        """
+        return getsizeof(self._to_dict())
+
+    @classmethod
+    def _from_dict(cls, input_dict: dict[str, Any]) -> RowMutationEntry:
+        """
+        Create a `RowMutationEntry` instance from a dictionary representation.
+
+        Args:
+            input_dict: A dictionary representation of the mutation entry.
+
+        Returns:
+            RowMutationEntry: A RowMutationEntry instance created from the dictionary.
+        """
+        return RowMutationEntry(
+            row_key=input_dict["row_key"],
+            mutations=[
+                Mutation._from_dict(mutation) for mutation in input_dict["mutations"]
+            ],
+        )
+
+
+@dataclass
+class _EntryWithProto:
+    """
+    A dataclass to hold a RowMutationEntry and its corresponding proto representation.
+
+    Used in _MutateRowsOperation to avoid repeated conversion of RowMutationEntry to proto.
+    """
+
+    entry: RowMutationEntry
+    proto: types_pb.MutateRowsRequest.Entry

google/cloud/bigtable/data/read_modify_write_rules.py

@@ -0,0 +1,112 @@
+# Copyright 2023 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 __future__ import annotations
+
+import abc
+
+import google.cloud.bigtable_v2.types.data as data_pb
+
+# value must fit in 64-bit signed integer
+_MAX_INCREMENT_VALUE = (1 << 63) - 1
+
+
+class ReadModifyWriteRule(abc.ABC):
+    """
+    Abstract base class for read-modify-write rules.
+    """
+
+    def __init__(self, family: str, qualifier: bytes | str):
+        qualifier = (
+            qualifier if isinstance(qualifier, bytes) else qualifier.encode("utf-8")
+        )
+        self.family = family
+        self.qualifier = qualifier
+
+    @abc.abstractmethod
+    def _to_dict(self) -> dict[str, str | bytes | int]:
+        raise NotImplementedError
+
+    def _to_pb(self) -> data_pb.ReadModifyWriteRule:
+        return data_pb.ReadModifyWriteRule(**self._to_dict())
+
+
+class IncrementRule(ReadModifyWriteRule):
+    """
+    Rule to increment a cell's value.
+
+    Args:
+        family:
+            The family name of the cell to increment.
+        qualifier:
+            The qualifier of the cell to increment.
+        increment_amount:
+            The amount to increment the cell's value. Must be between -2**63 and 2**63 (64-bit signed int).
+    Raises:
+        TypeError:
+            If increment_amount is not an integer.
+        ValueError:
+            If increment_amount is not between -2**63 and 2**63 (64-bit signed int).
+    """
+
+    def __init__(self, family: str, qualifier: bytes | str, increment_amount: int = 1):
+        if not isinstance(increment_amount, int):
+            raise TypeError("increment_amount must be an integer")
+        if abs(increment_amount) > _MAX_INCREMENT_VALUE:
+            raise ValueError(
+                "increment_amount must be between -2**63 and 2**63 (64-bit signed int)"
+            )
+        super().__init__(family, qualifier)
+        self.increment_amount = increment_amount
+
+    def _to_dict(self) -> dict[str, str | bytes | int]:
+        return {
+            "family_name": self.family,
+            "column_qualifier": self.qualifier,
+            "increment_amount": self.increment_amount,
+        }
+
+
+class AppendValueRule(ReadModifyWriteRule):
+    """
+    Rule to append a value to a cell's value.
+
+    Args:
+        family:
+            The family name of the cell to append to.
+        qualifier:
+            The qualifier of the cell to append to.
+        append_value:
+            The value to append to the cell's value.
+    Raises:
+        TypeError: If append_value is not bytes or str.
+    """
+
+    def __init__(self, family: str, qualifier: bytes | str, append_value: bytes | str):
+        append_value = (
+            append_value.encode("utf-8")
+            if isinstance(append_value, str)
+            else append_value
+        )
+        if not isinstance(append_value, bytes):
+            raise TypeError("append_value must be bytes or str")
+        super().__init__(family, qualifier)
+        self.append_value = append_value
+
+    def _to_dict(self) -> dict[str, str | bytes | int]:
+        return {
+            "family_name": self.family,
+            "column_qualifier": self.qualifier,
+            "append_value": self.append_value,
+        }