feat: add support for AddToCell in Data Client (#1147)

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,75 @@ def _to_dict(self) -> dict[str, Any]:
         }
 
 
+@dataclass
+class AddToCell(Mutation):
+    """
+    Adds an int64 value to an aggregate cell. The column family must be an
+    aggregate family and have an "int64" input type or this mutation will be
+    rejected.
+
+    Note: The timestamp values are in microseconds but must match the
+    granularity of the table (defaults to `MILLIS`). Therefore, the given value
+    must be a multiple of 1000 (millisecond granularity). For example:
+    `1571902339435000`.
+
+    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. Must be provided for
+          cell aggregation to work correctly.
+
+
+    Raises:
+        TypeError: If `qualifier` is not `bytes` or `str`.
+        TypeError: If `value` is not `int`.
+        TypeError: If `timestamp_micros` is not `int`.
+        ValueError: If `value` is out of bounds for a 64-bit signed int.
+        ValueError: If `timestamp_micros` is less than 0.
+    """
+
+    def __init__(
+        self,
+        family: str,
+        qualifier: bytes | str,
+        value: int,
+        timestamp_micros: int,
+    ):
+        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 not isinstance(timestamp_micros, int):
+            raise TypeError("timestamp_micros 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 < 0:
+            raise ValueError("timestamp must be non-negative")
+
+        self.family = family
+        self.qualifier = qualifier
+        self.value = value
+        self.timestamp = 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},
+                "input": {"int_value": self.value},
+            }
+        }
+
+    def is_idempotent(self) -> bool:
+        return False
+
+
 class RowMutationEntry:
     """
     A single entry in a `MutateRows` request.

tests/system/data/__init__.py

@@ -16,3 +16,4 @@
 
 TEST_FAMILY = "test-family"
 TEST_FAMILY_2 = "test-family-2"
+TEST_AGGREGATE_FAMILY = "test-aggregate-family"

tests/system/data/setup_fixtures.py

@@ -20,7 +20,7 @@
 import os
 import uuid
 
-from . import TEST_FAMILY, TEST_FAMILY_2
+from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY
 
 # authorized view subset to allow all qualifiers
 ALLOW_ALL = ""
@@ -183,6 +183,7 @@ def authorized_view_id(
                         "family_subsets": {
                             TEST_FAMILY: ALL_QUALIFIERS,
                             TEST_FAMILY_2: ALL_QUALIFIERS,
+                            TEST_AGGREGATE_FAMILY: ALL_QUALIFIERS,
                         },
                     },
                 },

tests/system/data/test_system_async.py

@@ -27,7 +27,7 @@
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
-from . import TEST_FAMILY, TEST_FAMILY_2
+from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY
 
 
 __CROSS_SYNC_OUTPUT__ = "tests.system.data.test_system_autogen"
@@ -76,6 +76,27 @@ async def add_row(
         await self.target.client._gapic_client.mutate_row(request)
         self.rows.append(row_key)
 
+    @CrossSync.convert
+    async def add_aggregate_row(
+        self, row_key, *, family=TEST_AGGREGATE_FAMILY, qualifier=b"q", input=0
+    ):
+        request = {
+            "table_name": self.target.table_name,
+            "row_key": row_key,
+            "mutations": [
+                {
+                    "add_to_cell": {
+                        "family_name": family,
+                        "column_qualifier": {"raw_value": qualifier},
+                        "timestamp": {"raw_timestamp_micros": 0},
+                        "input": {"int_value": input},
+                    }
+                }
+            ],
+        }
+        await self.target.client._gapic_client.mutate_row(request)
+        self.rows.append(row_key)
+
     @CrossSync.convert
     async def delete_rows(self):
         if self.rows:
@@ -132,7 +153,17 @@ def column_family_config(self):
         """
         from google.cloud.bigtable_admin_v2 import types
 
-        return {TEST_FAMILY: types.ColumnFamily(), TEST_FAMILY_2: types.ColumnFamily()}
+        int_aggregate_type = types.Type.Aggregate(
+            input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}),
+            sum={},
+        )
+        return {
+            TEST_FAMILY: types.ColumnFamily(),
+            TEST_FAMILY_2: types.ColumnFamily(),
+            TEST_AGGREGATE_FAMILY: types.ColumnFamily(
+                value_type=types.Type(aggregate_type=int_aggregate_type)
+            ),
+        }
 
     @pytest.fixture(scope="session")
     def init_table_id(self):
@@ -281,6 +312,37 @@ async def test_mutation_set_cell(self, target, temp_rows):
         # ensure cell is updated
         assert (await self._retrieve_cell_value(target, row_key)) == new_value
 
+    @CrossSync.pytest
+    @pytest.mark.usefixtures("target")
+    @CrossSync.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    async def test_mutation_add_to_cell(self, target, temp_rows):
+        """
+        Test add to cell mutation
+        """
+        from google.cloud.bigtable.data.mutations import AddToCell
+
+        row_key = b"add_to_cell"
+        family = TEST_AGGREGATE_FAMILY
+        qualifier = b"test-qualifier"
+        # add row to temp_rows, for future deletion
+        await temp_rows.add_aggregate_row(row_key, family=family, qualifier=qualifier)
+        # set and check cell value
+        await target.mutate_row(
+            row_key, AddToCell(family, qualifier, 1, timestamp_micros=0)
+        )
+        encoded_result = await self._retrieve_cell_value(target, row_key)
+        int_result = int.from_bytes(encoded_result, byteorder="big")
+        assert int_result == 1
+        # update again
+        await target.mutate_row(
+            row_key, AddToCell(family, qualifier, 9, timestamp_micros=0)
+        )
+        encoded_result = await self._retrieve_cell_value(target, row_key)
+        int_result = int.from_bytes(encoded_result, byteorder="big")
+        assert int_result == 10
+
     @pytest.mark.skipif(
         bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't use splits"
     )
@@ -1123,7 +1185,7 @@ async def test_execute_query_simple(self, client, table_id, instance_id):
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
     async def test_execute_against_target(
-        self, client, instance_id, table_id, temp_rows
+        self, client, instance_id, table_id, temp_rows, column_family_config
     ):
         await temp_rows.add_row(b"row_key_1")
         result = await client.execute_query(
@@ -1138,14 +1200,19 @@ async def test_execute_against_target(
         assert family_map[b"q"] == b"test-value"
         assert len(rows[0][TEST_FAMILY_2]) == 0
         md = result.metadata
-        assert len(md) == 3
+        # we expect it to fetch each column family, plus _key
+        # add additional families here if column_family_config changes
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )
 
     @pytest.mark.skipif(
         bool(os.environ.get(BIGTABLE_EMULATOR)),
@@ -1248,7 +1315,7 @@ async def test_execute_query_params(self, client, table_id, instance_id):
         predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
     )
     async def test_execute_metadata_on_empty_response(
-        self, client, instance_id, table_id, temp_rows
+        self, client, instance_id, table_id, temp_rows, column_family_config
     ):
         await temp_rows.add_row(b"row_key_1")
         result = await client.execute_query(
@@ -1258,11 +1325,16 @@ async def test_execute_metadata_on_empty_response(
 
         assert len(rows) == 0
         md = result.metadata
-        assert len(md) == 3
+        # we expect it to fetch each column family, plus _key
+        # add additional families here if column_family_config change
+        assert len(md) == len(column_family_config) + 1
         assert md["_key"].column_type == SqlType.Bytes()
         assert md[TEST_FAMILY].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
         assert md[TEST_FAMILY_2].column_type == SqlType.Map(
             SqlType.Bytes(), SqlType.Bytes()
         )
+        assert md[TEST_AGGREGATE_FAMILY].column_type == SqlType.Map(
+            SqlType.Bytes(), SqlType.Int64()
+        )