fix(bigtable): validate DirectRow and table match in MutationsBatcher.mutate

packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py

@@ -25,6 +25,7 @@
     RetryExceptionGroup,
 )
 from google.cloud.bigtable.data.mutations import RowMutationEntry
+from google.cloud.bigtable.row import DirectRow
 
 FLUSH_COUNT = 100  # after this many elements, send out the batch
 
@@ -181,9 +182,21 @@ def mutate(self, row):
         :param row: :class:`~google.cloud.bigtable.row.DirectRow`.
 
         :raises: One of the following:
+            * :exc:`TypeError` if ``row`` is not a :class:`~google.cloud.bigtable.row.DirectRow`.
+            * :exc:`~.table.TableMismatchError` if ``row`` belongs to a different table.
             * :exc:`~.table._BigtableRetryableError` if any row returned a transient error.
             * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried
         """
+        if not isinstance(row, DirectRow):
+            raise TypeError(
+                "Bulk processing cannot be applied for conditional or append mutations."
+            )
+        if row.table is not None and row.table.name != self.table.name:
+            from google.cloud.bigtable.table import TableMismatchError
+
+            raise TableMismatchError(
+                f"Row {row.row_key} is a part of {row.table.name} table. Current table: {self.table.name}"
+            )
         self._batcher.append(RowMutationEntry(row.row_key, row._get_mutations()))
 
     def mutate_rows(self, rows):
@@ -201,6 +214,8 @@ def mutate_rows(self, rows):
         :param rows: list:[`~google.cloud.bigtable.row.DirectRow`].
 
         :raises: One of the following:
+            * :exc:`TypeError` if any row in ``rows`` is not a :class:`~google.cloud.bigtable.row.DirectRow`.
+            * :exc:`~.table.TableMismatchError` if any row in ``rows`` belongs to a different table.
             * :exc:`~.table._BigtableRetryableError` if any row returned a transient error.
             * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried
         """

packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py

@@ -369,3 +369,71 @@ def test_mutations_batcher_atexit_flushes_after_explicit_flush(_setup_batcher, r
 
         assert operation_mock.call_count == 2
         assert len(recwarn) == 0
+
+
+def test_mutations_batcher_mutate_w_bound_table_row(_setup_batcher):
+    table, operation_mock = _setup_batcher
+    with MutationsBatcher(table=table, flush_count=1) as mutation_batcher:
+        row = table.direct_row(b"row_key")
+        row.set_cell("cf1", b"c1", 1)
+        mutation_batcher.mutate(row)
+
+    operation_mock.assert_called_once()
+
+
+def test_mutations_batcher_mutate_wrong_row_type(_setup_batcher):
+    from google.cloud.bigtable.row import AppendRow, ConditionalRow
+    from google.cloud.bigtable.row_filters import PassAllFilter
+
+    table, operation_mock = _setup_batcher
+    cond_row = ConditionalRow(
+        row_key=b"cond_row_key", table=table, filter_=PassAllFilter(True)
+    )
+    cond_row.set_cell("cf1", b"c1", b"true_val", state=True)
+    cond_row.set_cell("cf1", b"c1", b"false_val", state=False)
+    append_row = AppendRow(row_key=b"append_row_key", table=table)
+    append_row.append_cell_value("cf1", b"c1", b"val")
+
+    with MutationsBatcher(table=table, flush_count=1) as mutation_batcher:
+        with pytest.raises(
+            TypeError,
+            match="Bulk processing cannot be applied for conditional or append mutations.",
+        ):
+            mutation_batcher.mutate(cond_row)
+
+        with pytest.raises(
+            TypeError,
+            match="Bulk processing cannot be applied for conditional or append mutations.",
+        ):
+            mutation_batcher.mutate(append_row)
+
+        with pytest.raises(
+            TypeError,
+            match="Bulk processing cannot be applied for conditional or append mutations.",
+        ):
+            mutation_batcher.mutate_rows([cond_row])
+
+    operation_mock.assert_not_called()
+
+
+def test_mutations_batcher_mutate_table_mismatch(_setup_batcher):
+    from google.cloud.bigtable.table import Table, TableMismatchError
+
+    table, operation_mock = _setup_batcher
+    other_table = Table("other-table", instance=table._instance)
+    row = DirectRow(row_key=b"row_key", table=other_table)
+    row.set_cell("cf1", b"c1", b"val")
+
+    expected_msg = (
+        rf"Row b'row_key' is a part of {other_table.name} table\. "
+        rf"Current table: {table.name}"
+    )
+
+    with MutationsBatcher(table=table, flush_count=1) as mutation_batcher:
+        with pytest.raises(TableMismatchError, match=expected_msg):
+            mutation_batcher.mutate(row)
+
+        with pytest.raises(TableMismatchError, match=expected_msg):
+            mutation_batcher.mutate_rows([row])
+
+    operation_mock.assert_not_called()