updated sync files

Author: daniel-sanche

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

@@ -84,7 +84,10 @@ 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,
         )
         self.timeout_generator = _attempt_timeout_generator(
             attempt_timeout, operation_timeout

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

@@ -19,6 +19,7 @@
 from __future__ import annotations
 from typing import Sequence, TYPE_CHECKING
 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
@@ -107,7 +108,10 @@ def start_operation(self) -> CrossSync._Sync_Impl.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._Sync_Impl.Iterable[Row]:
@@ -268,7 +272,7 @@ def merge_rows(
                             if self._operation_metric.active_attempt is not None:
                                 self._operation_metric.active_attempt.application_blocking_time_ns += (
                                     time.monotonic_ns() - block_time
-                                ) * 1000
+                                )
                             break
                         c = it.__next__()
                 except _ResetRow as e:
@@ -285,9 +289,10 @@ def merge_rows(
                     continue
                 except CrossSync._Sync_Impl.StopIteration:
                     raise InvalidChunk("premature end of stream")
+        except GeneratorExit as close_exception:
+            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)
             raise generic_exception
         else:
             self._operation_metric.end_with_success()

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

@@ -759,7 +759,6 @@ def __init__(
             default_retryable_errors or ()
         )
         self._metrics = BigtableClientSideMetricsController(
-            client._metrics_interceptor,
             handlers=[],
             project_id=self.client.project,
             instance_id=instance_id,
@@ -941,8 +940,9 @@ def read_row(
         )
         results_generator = row_merger.start_operation()
         try:
-            return results_generator.__next__()
-        except StopIteration:
+            results = [a for a in results_generator]
+            return results[0]
+        except IndexError:
             return None
 
     def read_rows_sharded(
@@ -1142,7 +1142,10 @@ 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,
             )
 
     def mutations_batcher(
@@ -1265,7 +1268,10 @@ 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,
             )
 
     def bulk_mutate_rows(
@@ -1371,7 +1377,7 @@ def check_and_mutate_row(
         ):
             false_case_mutations = [false_case_mutations]
         false_case_list = [m._to_pb() for m in false_case_mutations or []]
-        with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
+        with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op:
             result = self.client._gapic_client.check_and_mutate_row(
                 request=CheckAndMutateRowRequest(
                     true_mutations=true_case_list,
@@ -1425,7 +1431,7 @@ def read_modify_write_row(
             rules = [rules]
         if not rules:
             raise ValueError("rules must contain at least one item")
-        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
+        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op:
             result = self.client._gapic_client.read_modify_write_row(
                 request=ReadModifyWriteRowRequest(
                     rules=[rule._to_pb() for rule in rules],

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

@@ -15,13 +15,12 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
+from typing import Sequence
 import time
 from functools import wraps
-from google.cloud.bigtable.data._metrics.data_model import (
-    OPERATION_INTERCEPTOR_METADATA_KEY,
-)
 from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
 from google.cloud.bigtable.data._metrics.data_model import OperationState
+from google.cloud.bigtable.data._metrics.data_model import OperationType
 from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
 from grpc import UnaryUnaryClientInterceptor
 from grpc import UnaryStreamClientInterceptor
@@ -33,19 +32,7 @@ def _with_operation_from_metadata(func):
 
     @wraps(func)
     def wrapper(self, continuation, client_call_details, request):
-        found_operation_id: str | None = None
-        try:
-            new_metadata: list[tuple[str, str]] = []
-            if client_call_details.metadata:
-                for k, v in client_call_details.metadata:
-                    if k == OPERATION_INTERCEPTOR_METADATA_KEY:
-                        found_operation_id = v
-                    else:
-                        new_metadata.append((k, v))
-            client_call_details.metadata = new_metadata
-        except Exception:
-            pass
-        operation: "ActiveOperationMetric" = self.operation_map.get(found_operation_id)
+        operation: "ActiveOperationMetric" | None = ActiveOperationMetric.get_active()
         if operation:
             if (
                 operation.state == OperationState.CREATED
@@ -59,18 +46,13 @@ 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:
-        operation.end_attempt_with_status(exc)
-
-
-def _get_metadata(source):
+def _get_metadata(source) -> dict[str, str | bytes] | None:
     """Helper to extract metadata from a call or RpcError"""
     try:
-        return (source.trailing_metadata() or []) + (source.initial_metadata() or [])
+        metadata: Sequence[tuple[str.str | bytes]] = (
+            source.trailing_metadata() + source.initial_metadata()
+        )
+        return {k: v for (k, v) in metadata}
     except Exception:
         return None
 
@@ -82,53 +64,31 @@ class BigtableMetricsInterceptor(
     An async gRPC interceptor to add client metadata and print server metadata.
     """
 
-    def __init__(self):
-        super().__init__()
-        self.operation_map = {}
-
-    def register_operation(self, operation):
-        """Register an operation object to be tracked my the interceptor
-
-        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"""
-        self.operation_map[operation.uuid] = operation
-        operation.handlers.append(self)
-
-    def on_operation_complete(self, op):
-        if op.uuid in self.operation_map:
-            del self.operation_map[op.uuid]
-
-    def on_operation_cancelled(self, op):
-        self.on_operation_complete(op)
-
     @_with_operation_from_metadata
     def intercept_unary_unary(
         self, operation, continuation, client_call_details, request
     ):
-        encountered_exc: Exception | None = None
         metadata = None
         try:
             call = continuation(client_call_details, request)
             metadata = _get_metadata(call)
             return call
         except Exception as rpc_error:
             metadata = _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)
 
     @_with_operation_from_metadata
     def intercept_unary_stream(
         self, operation, continuation, client_call_details, request
     ):
         def response_wrapper(call):
-            has_first_response = operation.first_response_latency is not None
+            has_first_response = (
+                operation.first_response_latency_ns is not None
+                or operation.op_type != OperationType.READ_ROWS
+            )
             encountered_exc = None
             try:
                 for response in call:
@@ -143,10 +103,14 @@ def response_wrapper(call):
                 raise
             finally:
                 if call is not None:
-                    _end_attempt(operation, encountered_exc, _get_metadata(call))
+                    metadata = _get_metadata(encountered_exc or call)
+                    if metadata is not None:
+                        operation.add_response_metadata(metadata)
 
         try:
             return response_wrapper(continuation(client_call_details, request))
         except Exception as rpc_error:
-            _end_attempt(operation, rpc_error, _get_metadata(rpc_error))
+            metadata = _get_metadata(rpc_error)
+            if metadata is not None:
+                operation.add_response_metadata(metadata)
             raise rpc_error

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

@@ -151,7 +151,7 @@ def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]):
                         )
             yield mutations[start_idx:end_idx]
 
-    async def add_to_flow_with_metrics(
+    def add_to_flow_with_metrics(
         self,
         mutations: RowMutationEntry | list[RowMutationEntry],
         metrics_controller: BigtableClientSideMetricsController,
@@ -161,9 +161,8 @@ async def add_to_flow_with_metrics(
             metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS)
             flow_start_time = time.monotonic_ns()
             try:
-                value = await inner_generator.__anext__()
-            except StopAsyncIteration:
-                metric.cancel()
+                value = inner_generator.__next__()
+            except CrossSync._Sync_Impl.StopIteration:
                 return
             metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time
             yield (value, metric)