feat: read_modify_write and check_and_mutate_row

google/cloud/bigtable/client.py

@@ -55,15 +55,15 @@
 from google.cloud.bigtable._helpers import _make_metadata
 from google.cloud.bigtable._helpers import _convert_retry_deadline
 
+from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
+from google.cloud.bigtable.row_filters import RowFilter
 from google.cloud.bigtable.row_filters import StripValueTransformerFilter
 from google.cloud.bigtable.row_filters import CellsRowLimitFilter
 from google.cloud.bigtable.row_filters import RowFilterChain
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.mutations_batcher import MutationsBatcher
     from google.cloud.bigtable import RowKeySamples
-    from google.cloud.bigtable.row_filters import RowFilter
-    from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
 
 
 class BigtableDataClient(ClientWithProject):
@@ -770,10 +770,11 @@ async def bulk_mutate_rows(
     async def check_and_mutate_row(
         self,
         row_key: str | bytes,
-        predicate: RowFilter | None,
+        predicate: RowFilter | dict[str, Any] | None,
+        *,
         true_case_mutations: Mutation | list[Mutation] | None = None,
         false_case_mutations: Mutation | list[Mutation] | None = None,
-        operation_timeout: int | float | None = 60,
+        operation_timeout: int | float | None = 20,
     ) -> bool:
         """
         Mutates a row atomically based on the output of a predicate filter
@@ -807,17 +808,43 @@ async def check_and_mutate_row(
         Raises:
             - GoogleAPIError exceptions from grpc call
         """
-        raise NotImplementedError
+        operation_timeout = operation_timeout or self.default_operation_timeout
+        if operation_timeout <= 0:
+            raise ValueError("operation_timeout must be greater than 0")
+        row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key
+        if true_case_mutations is not None and not isinstance(
+            true_case_mutations, list
+        ):
+            true_case_mutations = [true_case_mutations]
+        true_case_dict = [m._to_dict() for m in true_case_mutations or []]
+        if false_case_mutations is not None and not isinstance(
+            false_case_mutations, list
+        ):
+            false_case_mutations = [false_case_mutations]
+        false_case_dict = [m._to_dict() for m in false_case_mutations or []]
+        if predicate is not None and not isinstance(predicate, dict):
+            predicate = predicate.to_dict()
+        metadata = _make_metadata(self.table_name, self.app_profile_id)
+        result = await self.client._gapic_client.check_and_mutate_row(
+            request={
+                "predicate_filter": predicate,
+                "true_mutations": true_case_dict,
+                "false_mutations": false_case_dict,
+                "table_name": self.table_name,
+                "row_key": row_key,
+                "app_profile_id": self.app_profile_id,
+            },
+            metadata=metadata,
+            timeout=operation_timeout,
+        )
+        return result.predicate_matched
 
     async def read_modify_write_row(
         self,
         row_key: str | bytes,
-        rules: ReadModifyWriteRule
-        | list[ReadModifyWriteRule]
-        | dict[str, Any]
-        | list[dict[str, Any]],
+        rules: ReadModifyWriteRule | list[ReadModifyWriteRule],
         *,
-        operation_timeout: int | float | None = 60,
+        operation_timeout: int | float | None = 20,
     ) -> Row:
         """
         Reads and modifies a row atomically according to input ReadModifyWriteRules,
@@ -841,7 +868,29 @@ async def read_modify_write_row(
         Raises:
             - GoogleAPIError exceptions from grpc call
         """
-        raise NotImplementedError
+        operation_timeout = operation_timeout or self.default_operation_timeout
+        row_key = row_key.encode("utf-8") if isinstance(row_key, str) else row_key
+        if operation_timeout <= 0:
+            raise ValueError("operation_timeout must be greater than 0")
+        if rules is not None and not isinstance(rules, list):
+            rules = [rules]
+        if not rules:
+            raise ValueError("rules must contain at least one item")
+        # concert to dict representation
+        rules_dict = [rule._to_dict() for rule in rules]
+        metadata = _make_metadata(self.table_name, self.app_profile_id)
+        result = await self.client._gapic_client.read_modify_write_row(
+            request={
+                "rules": rules_dict,
+                "table_name": self.table_name,
+                "row_key": row_key,
+                "app_profile_id": self.app_profile_id,
+            },
+            metadata=metadata,
+            timeout=operation_timeout,
+        )
+        # construct Row from result
+        return Row._from_pb(result.row)
 
     async def close(self):
         """

google/cloud/bigtable/mutations.py

@@ -18,6 +18,8 @@
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
 
+from google.cloud.bigtable.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
 
@@ -99,6 +101,10 @@ def __init__(
         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")

google/cloud/bigtable/read_modify_write_rules.py

@@ -14,22 +14,59 @@
 #
 from __future__ import annotations
 
-from dataclasses import dataclass
+import abc
 
+# value must fit in 64-bit signed integer
+MAX_INCREMENT_VALUE = (1 << 63) - 1
 
-class ReadModifyWriteRule:
-    pass
+
+class ReadModifyWriteRule(abc.ABC):
+    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):
+        raise NotImplementedError
 
 
-@dataclass
 class IncrementRule(ReadModifyWriteRule):
-    increment_amount: int
-    family: str
-    qualifier: bytes
+    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):
+        return {
+            "family_name": self.family,
+            "column_qualifier": self.qualifier,
+            "increment_amount": self.increment_amount,
+        }
 
 
-@dataclass
 class AppendValueRule(ReadModifyWriteRule):
-    append_value: bytes
-    family: str
-    qualifier: bytes
+    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):
+        return {
+            "family_name": self.family,
+            "column_qualifier": self.qualifier,
+            "append_value": self.append_value,
+        }

google/cloud/bigtable/row.py

@@ -18,6 +18,8 @@
 from typing import Sequence, Generator, overload, Any
 from functools import total_ordering
 
+from google.cloud.bigtable_v2.types import Row as RowPB
+
 # Type aliases used internally for readability.
 _family_type = str
 _qualifier_type = bytes
@@ -72,6 +74,30 @@ def _index(
                 ).append(cell)
         return self._index_data
 
+    @classmethod
+    def _from_pb(cls, row_pb: RowPB) -> Row:
+        """
+        Creates a row from a protobuf representation
+
+        Row objects are not intended to be created by users.
+        They are returned by the Bigtable backend.
+        """
+        row_key: bytes = row_pb.key
+        cell_list: list[Cell] = []
+        for family in row_pb.families:
+            for column in family.columns:
+                for cell in column.cells:
+                    new_cell = Cell(
+                        value=cell.value,
+                        row_key=row_key,
+                        family=family.name,
+                        qualifier=column.qualifier,
+                        timestamp_micros=cell.timestamp_micros,
+                        labels=list(cell.labels) if cell.labels else None,
+                    )
+                    cell_list.append(new_cell)
+        return cls(row_key, cells=cell_list)
+
     def get_cells(
         self, family: str | None = None, qualifier: str | bytes | None = None
     ) -> list[Cell]: