added bulk_mutate_row tests

Author: daniel-sanche

google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -16,6 +16,8 @@
 
 from typing import Sequence, TYPE_CHECKING
 
+from functools import partial
+
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry as retries
 import google.cloud.bigtable_v2.types.bigtable as types_pb
@@ -106,7 +108,9 @@ def __init__(
             self.is_retryable,
             metric.backoff_generator,
             operation_timeout,
-            exception_factory=_retry_exception_factory,
+            exception_factory=partial(
+                _retry_exception_factory, operation=metric
+            ),
         )
         # initialize state
         self.timeout_generator = _attempt_timeout_generator(

google/cloud/bigtable/data/_metrics/data_model.py

@@ -418,8 +418,11 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
         The operation is automatically ended on exit, with the status determined
         by the exception type and value.
+
+        If operation was already ended manually, do nothing.
         """
-        if exc_val is None:
-            self.end_with_success()
-        else:
-            self.end_with_status(exc_val)
+        if not self.state == OperationState.COMPLETED:
+            if exc_val is None:
+                self.end_with_success()
+            else:
+                self.end_with_status(exc_val)

tests/system/data/test_metrics_async.py

@@ -830,23 +830,35 @@ async def test_read_rows_sharded_failure_mid_stream(
         query2 = ReadRowsQuery(row_keys=[b"b"])
         handler.clear()
         error_injector.fail_mid_stream = True
-        error_injector.push(PermissionDenied("terminal"))
+        error_injector.push(Aborted("retryable"))
         error_injector.push(PermissionDenied("terminal"))
         with pytest.raises(ShardedReadRowsExceptionGroup) as e:
-            await table.read_rows_sharded([query1, query2])
-        assert len(e.value.exceptions) == 2
+            await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted])
+        assert len(e.value.exceptions) == 1
         assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied)
+        # one shard will fail, the other will succeed
+        # the failing shard will have one retry
         assert len(handler.completed_operations) == 2
-        assert len(handler.completed_attempts) == 2
+        assert len(handler.completed_attempts) == 3
         assert len(handler.cancelled_operations) == 0
-        for operation in handler.completed_operations:
-            assert operation.final_status.name == "PERMISSION_DENIED"
-            assert operation.op_type.value == "ReadRows"
-            assert operation.is_streaming is True
-            assert len(operation.completed_attempts) == 1
-            # validate attempt
-            attempt = operation.completed_attempts[0]
-            assert attempt.end_status.name == "PERMISSION_DENIED"
+        # sort operations by status
+        failed_op = next(op for op in handler.completed_operations if op.final_status.name != "OK")
+        success_op = next(op for op in handler.completed_operations if op.final_status.name == "OK")
+        # validate failed operation
+        assert failed_op.final_status.name == "PERMISSION_DENIED"
+        assert failed_op.op_type.value == "ReadRows"
+        assert failed_op.is_streaming is True
+        assert len(failed_op.completed_attempts) == 2
+        # validate retried attempt
+        attempt = failed_op.completed_attempts[0]
+        assert attempt.end_status.name == "ABORTED"
+        # validate final attempt
+        final_attempt = failed_op.completed_attempts[-1]
+        assert final_attempt.end_status.name == "PERMISSION_DENIED"
+        # validate successful operation
+        assert success_op.final_status.name == "OK"
+        assert len(success_op.completed_attempts) == 1
+        assert success_op.completed_attempts[0].end_status.name == "OK"
 
     @CrossSync.pytest
     async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config):
@@ -896,7 +908,45 @@ async def test_bulk_mutate_rows_failure_grpc(
 
         No headers expected
         """
-        pass
+        from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell
+        from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+
+        row_key = b"row_key_1"
+        mutation = SetCell(TEST_FAMILY, b"q", b"v")
+        entry = RowMutationEntry(row_key, [mutation])
+        assert entry.is_idempotent()
+
+        handler.clear()
+        exc = Aborted("injected")
+        num_retryable = 2
+        for i in range(num_retryable):
+            error_injector.push(exc)
+        error_injector.push(PermissionDenied("terminal"))
+        with pytest.raises(MutationsExceptionGroup) as e:
+            await table.bulk_mutate_rows([entry], retryable_errors=[Aborted])
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) == num_retryable + 1
+        assert len(handler.cancelled_operations) == 0
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
+        assert operation.final_status.name == "PERMISSION_DENIED"
+        assert operation.op_type.value == "MutateRows"
+        assert operation.is_streaming is False
+        assert len(operation.completed_attempts) == num_retryable + 1
+        assert operation.cluster_id == "unspecified"
+        assert operation.zone == "global"
+        # validate attempts
+        for i in range(num_retryable):
+            attempt = handler.completed_attempts[i]
+            assert isinstance(attempt, CompletedAttemptMetric)
+            assert attempt.end_status.name == "ABORTED"
+            assert attempt.gfe_latency_ns is None
+        final_attempt = handler.completed_attempts[num_retryable]
+        assert isinstance(final_attempt, CompletedAttemptMetric)
+        assert final_attempt.end_status.name == "PERMISSION_DENIED"
+        assert final_attempt.gfe_latency_ns is None
 
     @CrossSync.pytest
     async def test_bulk_mutate_rows_failure_timeout(
@@ -907,7 +957,35 @@ async def test_bulk_mutate_rows_failure_timeout(
 
         No grpc headers expected
         """
-        pass
+        from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell
+        from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+        from google.api_core.exceptions import DeadlineExceeded
+
+        row_key = b"row_key_1"
+        mutation = SetCell(TEST_FAMILY, b"q", b"v")
+        entry = RowMutationEntry(row_key, [mutation])
+
+        handler.clear()
+        with pytest.raises(MutationsExceptionGroup) as e:
+            await table.bulk_mutate_rows([entry], operation_timeout=0.001)
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) == 1
+        assert len(handler.cancelled_operations) == 0
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
+        assert operation.final_status.name == "DEADLINE_EXCEEDED"
+        assert operation.op_type.value == "MutateRows"
+        assert operation.is_streaming is False
+        assert len(operation.completed_attempts) == 1
+        assert operation.cluster_id == "unspecified"
+        assert operation.zone == "global"
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert isinstance(attempt, CompletedAttemptMetric)
+        assert attempt.end_status.name == "DEADLINE_EXCEEDED"
+        assert attempt.gfe_latency_ns is None
 
     @CrossSync.pytest
     async def test_bulk_mutate_rows_failure_unauthorized(
@@ -916,7 +994,35 @@ async def test_bulk_mutate_rows_failure_unauthorized(
         """
         Test failure in backend by accessing an unauthorized family
         """
-        pass
+        from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell
+        from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+
+        row_key = b"row_key_1"
+        mutation = SetCell("unauthorized", b"q", b"v")
+        entry = RowMutationEntry(row_key, [mutation])
+
+        handler.clear()
+        with pytest.raises(MutationsExceptionGroup) as e:
+            await authorized_view.bulk_mutate_rows([entry])
+        assert len(e.value.exceptions) == 1
+        assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError)
+        assert e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED"
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) == 1
+        assert len(handler.cancelled_operations) == 0
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert operation.final_status.name == "PERMISSION_DENIED"
+        assert operation.op_type.value == "MutateRows"
+        assert operation.is_streaming is False
+        assert len(operation.completed_attempts) == 1
+        assert operation.cluster_id == next(iter(cluster_config.keys()))
+        assert operation.zone == cluster_config[operation.cluster_id].location.split("/")[-1]
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert attempt.end_status.name == "PERMISSION_DENIED"
+        assert attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns
 
     @CrossSync.pytest
     async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config):