feat: add support for AddToCell in Data Client

Author: axyjo

google/cloud/bigtable/data/mutations.py

@@ -123,6 +123,14 @@ def _from_dict(cls, input_dict: dict[str, Any]) -> Mutation:
                 instance = DeleteAllFromFamily(details["family_name"])
             elif "delete_from_row" in input_dict:
                 instance = DeleteAllFromRow()
+            elif "add_to_cell" in input_dict:
+                details = input_dict["add_to_cell"]
+                instance = AddToCell(
+                    details["family_name"],
+                    details["column_qualifier"]["raw_value"],
+                    details["input"]["int_value"],
+                    details["timestamp"]["raw_timestamp_micros"],
+                )
         except KeyError as e:
             raise ValueError("Invalid mutation dictionary") from e
         if instance is None:
@@ -276,6 +284,65 @@ def _to_dict(self) -> dict[str, Any]:
         }
 
 
+@dataclass
+class AddToCell(Mutation):
+    """
+    Mutation to add a value to an aggregate cell.
+
+
+    Args:
+        family: The name of the column family to which the cell belongs.
+        qualifier: The column qualifier of the cell.
+        value: The value to be accumulated into the cell.
+        timestamp_micros: The timestamp of the cell.
+
+    Raises:
+        TypeError: If `qualifier` is not `bytes` or `str`.
+        TypeError: If `value` is not `int`.
+        ValueError: If `timestamp_micros` is less than `_SERVER_SIDE_TIMESTAMP`.
+    """
+
+    def __init__(
+        self,
+        family: str,
+        qualifier: bytes | str,
+        value: 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 not isinstance(value, int):
+            raise TypeError("value must be int")
+        if abs(value) > _MAX_INCREMENT_VALUE:
+            raise ValueError(
+                "int values must be between -2**63 and 2**63 (64-bit signed 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)
+
+        self.family = family
+        self.qualifier = qualifier
+        self.value = value
+        self.timestamp_micros = timestamp_micros
+
+    def _to_dict(self) -> dict[str, Any]:
+        return {
+            "add_to_cell": {
+                "family_name": self.family,
+                "column_qualifier": {"raw_value": self.qualifier},
+                "timestamp": {"raw_timestamp_micros": self.timestamp_micros},
+                "input": {"int_value": self.value},
+            }
+        }
+
+    def is_idempotent(self) -> bool:
+        return False
+
+
 class RowMutationEntry:
     """
     A single entry in a `MutateRows` request.

tests/unit/data/test_mutations.py

@@ -706,3 +706,131 @@ def test__from_dict(self):
         assert len(instance.mutations) == 1
         assert isinstance(instance.mutations[0], mutations.DeleteAllFromFamily)
         assert instance.mutations[0].family_to_delete == "test_family"
+
+
+class TestAddToCell:
+    def _target_class(self):
+        from google.cloud.bigtable.data.mutations import AddToCell
+
+        return AddToCell
+
+    def _make_one(self, *args, **kwargs):
+        return self._target_class()(*args, **kwargs)
+
+    @pytest.mark.parametrize("input_val", [2**64, -(2**64)])
+    def test_ctor_large_int(self, input_val):
+        with pytest.raises(ValueError) as e:
+            self._make_one(family="f", qualifier=b"b", value=input_val)
+        assert "int values must be between" in str(e.value)
+
+    @pytest.mark.parametrize("input_val", ["", "a", "abc", "hello world!"])
+    def test_ctor_str_value(self, input_val):
+        with pytest.raises(ValueError) as e:
+            self._make_one(family="f", qualifier=b"b", value=input_val)
+        assert "value must be int" in str(e.value)
+
+    def test_ctor(self):
+        """Ensure constructor sets expected values"""
+        expected_family = "test-family"
+        expected_qualifier = b"test-qualifier"
+        expected_value = 1234
+        expected_timestamp = 1234567890
+        instance = self._make_one(
+            expected_family, expected_qualifier, expected_value, expected_timestamp
+        )
+        assert instance.family == expected_family
+        assert instance.qualifier == expected_qualifier
+        assert instance.value == expected_value
+        assert instance.timestamp_micros == expected_timestamp
+
+    def test_ctor_negative_timestamp(self):
+        """Only positive or -1 timestamps are valid"""
+        with pytest.raises(ValueError) as e:
+            self._make_one("test-family", b"test-qualifier", b"test-value", -2)
+        assert (
+            "timestamp_micros must be positive (or -1 for server-side timestamp)"
+            in str(e.value)
+        )
+
+    @pytest.mark.parametrize(
+        "timestamp_ns,expected_timestamp_micros",
+        [
+            (0, 0),
+            (1, 0),
+            (123, 0),
+            (999, 0),
+            (999_999, 0),
+            (1_000_000, 1000),
+            (1_234_567, 1000),
+            (1_999_999, 1000),
+            (2_000_000, 2000),
+            (1_234_567_890_123, 1_234_567_000),
+        ],
+    )
+    def test_ctor_no_timestamp(self, timestamp_ns, expected_timestamp_micros):
+        """If no timestamp is given, should use current time with millisecond precision"""
+        with mock.patch("time.time_ns", return_value=timestamp_ns):
+            instance = self._make_one("test-family", b"test-qualifier", 1234)
+            assert instance.timestamp_micros == expected_timestamp_micros
+
+    def test__to_dict(self):
+        """ensure dict representation is as expected"""
+        expected_family = "test-family"
+        expected_qualifier = b"test-qualifier"
+        expected_value = 1234
+        expected_timestamp = 123456789
+        instance = self._make_one(
+            expected_family, expected_qualifier, expected_value, expected_timestamp
+        )
+        got_dict = instance._to_dict()
+        assert list(got_dict.keys()) == ["add_to_cell"]
+        got_inner_dict = got_dict["add_to_cell"]
+        assert got_inner_dict["family_name"] == expected_family
+        assert got_inner_dict["column_qualifier"]["raw_value"] == expected_qualifier
+        assert got_inner_dict["timestamp_micros"]["raw_timestamp_micros"] == expected_timestamp
+        assert got_inner_dict["value"]["int_value"] == expected_value
+        assert len(got_inner_dict.keys()) == 4
+
+    def test__to_pb(self):
+        """ensure proto representation is as expected"""
+        import google.cloud.bigtable_v2.types.data as data_pb
+
+        expected_family = "test-family"
+        expected_qualifier = b"test-qualifier"
+        expected_value = 1234
+        expected_timestamp = 123456789
+        instance = self._make_one(
+            expected_family, expected_qualifier, expected_value, expected_timestamp
+        )
+        got_pb = instance._to_pb()
+        assert isinstance(got_pb, data_pb.Mutation)
+        assert got_pb.set_cell.family_name == expected_family
+        assert got_pb.set_cell.column_qualifier.raw_value == expected_qualifier
+        assert got_pb.set_cell.timestamp_micros.raw_timestamp_micros == expected_timestamp
+        assert got_pb.set_cell.value.int_value == expected_value
+
+    @pytest.mark.parametrize(
+        "timestamp",
+        [
+            (1234567890),
+            (1),
+            (0),
+            (-1),
+            (None),
+        ],
+    )
+    def test_is_idempotent(self, timestamp, expected_value):
+        """is_idempotent is not based on whether an explicit timestamp is set"""
+        instance = self._make_one(
+            "test-family", b"test-qualifier", 1234, timestamp
+        )
+        assert not instance.is_idempotent()
+
+    def test___str__(self):
+        """Str representation of mutations should be to_dict"""
+        instance = self._make_one(
+            "test-family", b"test-qualifier", 1234, 1234567890
+        )
+        str_value = instance.__str__()
+        dict_value = instance._to_dict()
+        assert str_value == str(dict_value)