added metric system tests; solved issues

Author: daniel-sanche

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

@@ -106,7 +106,8 @@ def __init__(
             self.is_retryable,
             metric.backoff_generator,
             operation_timeout,
-            exception_factory=_retry_exception_factory,
+            exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
+            on_error=self._operation_metric.track_retryable_error,
         )
         # initialize state
         self.timeout_generator = _attempt_timeout_generator(

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

@@ -19,6 +19,8 @@
 
 import time
 
+from grpc import StatusCode
+
 from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
 from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB
 from google.cloud.bigtable_v2.types import RowSet as RowSetPB
@@ -121,7 +123,8 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
             self._predicate,
             self._operation_metric.backoff_generator,
             self.operation_timeout,
-            exception_factory=_retry_exception_factory,
+            exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
+            on_error=self._operation_metric.track_retryable_error,
         )
 
     def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
@@ -322,7 +325,7 @@ async def merge_rows(
                             if self._operation_metric.active_attempt is not None:
                                 self._operation_metric.active_attempt.application_blocking_time_ns += (  # type: ignore
                                     time.monotonic_ns() - block_time
-                                ) * 1000
+                                )
                             break
                         c = await it.__anext__()
                 except _ResetRow as e:
@@ -339,9 +342,12 @@ async def merge_rows(
                     continue
                 except CrossSync.StopIteration:
                     raise InvalidChunk("premature end of stream")
+        except GeneratorExit as close_exception:
+            # handle aclose()
+            self._operation_metric.end_with_status(StatusCode.CANCELLED)
+            raise close_exception
         except Exception as generic_exception:
-            if not self._predicate(generic_exception):
-                self._operation_metric.end_attempt_with_status(generic_exception)
+            # handle exceptions in retry wrapper
             raise generic_exception
         else:
             self._operation_metric.end_with_success()

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

@@ -1391,7 +1391,8 @@ async def execute_rpc():
                 predicate,
                 operation_metric.backoff_generator,
                 operation_timeout,
-                exception_factory=_retry_exception_factory,
+                exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
+                on_error=operation_metric.track_retryable_error,
             )
 
     @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"})
@@ -1525,7 +1526,8 @@ async def mutate_row(
                 predicate,
                 operation_metric.backoff_generator,
                 operation_timeout,
-                exception_factory=_retry_exception_factory,
+                exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
+                on_error=operation_metric.track_retryable_error,
             )
 
     @CrossSync.convert

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

@@ -13,8 +13,12 @@
 # limitations under the License
 from __future__ import annotations
 
+from typing import Sequence
+
 import time
 from functools import wraps
+from grpc import StatusCode
+
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -28,6 +32,7 @@
 if CrossSync.is_async:
     from grpc.aio import UnaryUnaryClientInterceptor
     from grpc.aio import UnaryStreamClientInterceptor
+    from grpc.aio import AioRpcError
 else:
     from grpc import UnaryUnaryClientInterceptor
     from grpc import UnaryStreamClientInterceptor
@@ -75,23 +80,21 @@ def wrapper(self, continuation, client_call_details, request):
 
     return wrapper
 
-
-def _end_attempt(operation, exc, metadata):
-    """Helper to add metadata and exception to an operation"""
-    if metadata is not None:
-        operation.add_response_metadata(metadata)
-    if exc is not None:
-        # end attempt. If it succeeded, let higher levels decide when to end operation
-        operation.end_attempt_with_status(exc)
-
-
 @CrossSync.convert
-async def _get_metadata(source):
+async def _get_metadata(source) -> dict[str, str|bytes] | None:
     """Helper to extract metadata from a call or RpcError"""
     try:
-        return (await source.trailing_metadata() or []) + (
-            await source.initial_metadata() or []
-        )
+        if CrossSync.is_async:
+            # grpc.aio returns metadata in Metadata objects
+            if isinstance(source, AioRpcError):
+                metadata = list(source.trailing_metadata()) + list(source.initial_metadata())
+            else:
+                metadata = list(await source.trailing_metadata()) + list(await source.initial_metadata())
+        else:
+            # sync grpc returns metadata as a sequence of tuples
+            metadata: Sequence[tuple[str. str|bytes]] = source.trailing_metadata() + source.initial_metadata()
+        # convert metadata to dict format
+        return {k:v for k,v in metadata}
     except Exception:
         # ignore errors while fetching metadata
         return None
@@ -116,7 +119,6 @@ def register_operation(self, operation):
         When registered, the operation will receive metadata updates:
         - start_attempt if attempt not started when rpc is being sent
         - add_response_metadata after call is complete
-        - end_attempt_with_status if attempt receives an error
 
         The interceptor will register itself as a handeler for the operation,
         so it can unregister the operation when it is complete
@@ -136,18 +138,17 @@ def on_operation_cancelled(self, op):
     async def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
-        encountered_exc: Exception | None = None
         metadata = None
         try:
             call = await continuation(client_call_details, request)
             metadata = await _get_metadata(call)
             return call
         except Exception as rpc_error:
             metadata = await _get_metadata(rpc_error)
-            encountered_exc = rpc_error
             raise rpc_error
         finally:
-            _end_attempt(operation, encountered_exc, metadata)
+            if metadata is not None:
+                operation.add_response_metadata(metadata)
 
     @CrossSync.convert
     @_with_operation_from_metadata
@@ -156,7 +157,10 @@ async def intercept_unary_stream(
     ):
         async def response_wrapper(call):
             # only track has_first response for READ_ROWS
-            has_first_response = operation.first_response_latency_ns is not None or operation.op_type != OperationType.READ_ROWS
+            has_first_response = (
+                operation.first_response_latency_ns is not None
+                or operation.op_type != OperationType.READ_ROWS
+            )
             encountered_exc = None
             try:
                 async for response in call:
@@ -173,11 +177,14 @@ async def response_wrapper(call):
                 raise
             finally:
                 if call is not None:
-                    _end_attempt(operation, encountered_exc, await _get_metadata(call))
+                    metadata = await _get_metadata(encountered_exc or call)
+                    if metadata is not None:
+                        operation.add_response_metadata(metadata)
 
         try:
             return response_wrapper(await continuation(client_call_details, request))
         except Exception as rpc_error:
-            # handle errors while intializing stream
-            _end_attempt(operation, rpc_error, await _get_metadata(rpc_error))
+            metadata = await _get_metadata(rpc_error)
+            if metadata is not None:
+                operation.add_response_metadata(metadata)
             raise rpc_error

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

@@ -181,6 +181,7 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]
                         )
             yield mutations[start_idx:end_idx]
 
+    @CrossSync.convert(replace_symbols={"__anext__": "__next__"})
     async def add_to_flow_with_metrics(
         self,
         mutations: RowMutationEntry | list[RowMutationEntry],
@@ -193,7 +194,7 @@ async def add_to_flow_with_metrics(
             flow_start_time = time.monotonic_ns()
             try:
                 value = await inner_generator.__anext__()
-            except StopAsyncIteration:
+            except CrossSync.StopIteration:
                 metric.cancel()
                 return
             metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time

google/cloud/bigtable/data/_helpers.py

@@ -16,7 +16,7 @@
 """
 from __future__ import annotations
 
-from typing import Sequence, List, Tuple, TYPE_CHECKING, Union
+from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union
 import time
 import enum
 from collections import namedtuple
@@ -103,6 +103,7 @@ def _retry_exception_factory(
         tuple[Exception, Exception|None]:
             tuple of the exception to raise, and a cause exception if applicable
     """
+    exc_list = exc_list.copy()
     if reason == RetryFailureReason.TIMEOUT:
         timeout_val_str = f"of {timeout_val:0.1f}s " if timeout_val is not None else ""
         # if failed due to timeout, raise deadline exceeded as primary exception

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

@@ -25,10 +25,16 @@
 from dataclasses import dataclass
 from dataclasses import field
 from grpc import StatusCode
+from grpc import RpcError
+from grpc.aio import AioRpcError
 
+from google.api_core.exceptions import GoogleAPICallError
+from google.api_core.retry import RetryFailureReason
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable_v2.types.response_params import ResponseParams
 from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
+from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
+from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
 from google.protobuf.message import DecodeError
 
 if TYPE_CHECKING:
@@ -280,7 +286,7 @@ def _parse_response_metadata_blob(blob: bytes) -> Tuple[str, str] | None:
             # failed to parse metadata
             return None
 
-    def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
+    def end_attempt_with_status(self, status: StatusCode | BaseException) -> None:
         """
         Called to mark the end of an attempt for the operation.
 
@@ -297,7 +303,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
             return self._handle_error(
                 INVALID_STATE_ERROR.format("end_attempt_with_status", self.state)
             )
-        if isinstance(status, Exception):
+        if isinstance(status, BaseException):
             status = self._exc_to_status(status)
         complete_attempt = CompletedAttemptMetric(
             duration_ns=time.monotonic_ns() - self.active_attempt.start_time_ns,
@@ -312,7 +318,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None:
         for handler in self.handlers:
             handler.on_attempt_complete(complete_attempt, self)
 
-    def end_with_status(self, status: StatusCode | Exception) -> None:
+    def end_with_status(self, status: StatusCode | BaseException) -> None:
         """
         Called to mark the end of the operation. If there is an active attempt,
         end_attempt_with_status will be called with the same status.
@@ -329,7 +335,7 @@ def end_with_status(self, status: StatusCode | Exception) -> None:
                 INVALID_STATE_ERROR.format("end_with_status", self.state)
             )
         final_status = (
-            self._exc_to_status(status) if isinstance(status, Exception) else status
+            self._exc_to_status(status) if isinstance(status, BaseException) else status
         )
         if self.state == OperationState.ACTIVE_ATTEMPT:
             self.end_attempt_with_status(final_status)
@@ -367,7 +373,7 @@ def cancel(self):
             handler.on_operation_cancelled(self)
 
     @staticmethod
-    def _exc_to_status(exc: Exception) -> StatusCode:
+    def _exc_to_status(exc: BaseException) -> StatusCode:
         """
         Extracts the grpc status code from an exception.
 
@@ -389,8 +395,71 @@ def _exc_to_status(exc: Exception) -> StatusCode:
             and exc.__cause__.grpc_status_code is not None
         ):
             return exc.__cause__.grpc_status_code
+        if isinstance(exc, AioRpcError) or isinstance(exc, RpcError):
+            return exc.code()
         return StatusCode.UNKNOWN
 
+    def track_retryable_error(self, exc: Exception) -> None:
+        """
+        Used as input to api_core.Retry classes, to track when retryable errors are encountered
+
+        Should be passed as on_error callback
+        """
+        try:
+            # record metadata from failed rpc
+            if (
+                isinstance(exc, GoogleAPICallError)
+                and exc.errors
+            ):
+                rpc_error = exc.errors[-1]
+                metadata = list(rpc_error.trailing_metadata()) + list(
+                    rpc_error.initial_metadata()
+                )
+                self.add_response_metadata({k: v for k, v in metadata})
+        except Exception:
+            # ignore errors in metadata collection
+            pass
+        if isinstance(exc, _MutateRowsIncomplete):
+            # _MutateRowsIncomplete represents a successful rpc with some failed mutations
+            # mark the attempt as successful
+            self.end_attempt_with_status(StatusCode.OK)
+        else:
+            self.end_attempt_with_status(exc)
+
+    def track_terminal_error(self, exception_factory:callable[
+        [list[Exception], RetryFailureReason, float | None],tuple[Exception, Exception | None],
+    ]) -> callable[[list[Exception], RetryFailureReason, float | None], None]:
+        """
+        Used as input to api_core.Retry classes, to track when terminal errors are encountered
+
+        Should be used as a wrapper over an exception_factory callback
+        """
+        def wrapper(
+            exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None
+        ) -> tuple[Exception, Exception | None]:
+            source_exc, cause_exc = exception_factory(exc_list, reason, timeout_val)
+            try:
+                # record metadata from failed rpc
+                if (
+                    isinstance(source_exc, GoogleAPICallError)
+                    and source_exc.errors
+                ):
+                    rpc_error = source_exc.errors[-1]
+                    metadata = list(rpc_error.trailing_metadata()) + list(
+                        rpc_error.initial_metadata()
+                    )
+                    self.add_response_metadata({k: v for k, v in metadata})
+            except Exception:
+                # ignore errors in metadata collection
+                pass
+            if reason == RetryFailureReason.TIMEOUT and self.state == OperationState.ACTIVE_ATTEMPT and exc_list:
+                # record ending attempt for timeout failures
+                attempt_exc = exc_list[-1]
+                self.track_retryable_error(attempt_exc)
+            self.end_with_status(source_exc)
+            return source_exc, cause_exc
+        return wrapper
+
     @staticmethod
     def _handle_error(message: str) -> None:
         """
@@ -418,8 +487,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/__init__.py

@@ -15,12 +15,12 @@
 #
 import pytest
 import uuid
-from google.cloud.bigtable.data._cross_sync import CrossSync
 
 TEST_FAMILY = "test-family"
 TEST_FAMILY_2 = "test-family-2"
 TEST_AGGREGATE_FAMILY = "test-aggregate-family"
 
+
 class SystemTestRunner:
     """
     configures a system test class with configuration for clusters/tables/etc
@@ -67,4 +67,4 @@ def column_family_config(self):
             TEST_AGGREGATE_FAMILY: types.ColumnFamily(
                 value_type=types.Type(aggregate_type=int_aggregate_type)
             ),
-        }
+        }