diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index 72940c1ce050..14f857467929 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/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 """ diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py index b30c4fcff4e7..6c326dd40fb8 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py +++ b/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()