fixed mypy

Author: daniel-sanche

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

@@ -220,7 +220,7 @@ async def chunk_stream(
     )
     async def merge_rows(
         chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None,
-        attempt_metric: ActiveAttemptMetric,
+        attempt_metric: ActiveAttemptMetric | None,
     ) -> CrossSync.Iterable[Row]:
         """
         Merge chunks into rows
@@ -233,7 +233,6 @@ async def merge_rows(
         if chunks is None:
             return
         it = chunks.__aiter__()
-        is_first_row = True
         # For each row
         while True:
             try:
@@ -316,17 +315,14 @@ async def merge_rows(
                         Cell(value, row_key, family, qualifier, ts, list(labels))
                     )
                     if c.commit_row:
-                        if is_first_row:
-                            # record first row latency in metrics
-                            is_first_row = False
-                            attempt_metric.attempt_first_response()
                         block_time = time.monotonic_ns()
                         yield Row(row_key, cells)
                         # most metric operations use setters, but this one updates
                         # the value directly to avoid extra overhead
-                        attempt_metric.active_attempt.application_blocking_time_ns += (  # type: ignore
-                            time.monotonic_ns() - block_time
-                        ) * 1000
+                        if attempt_metric is not None:
+                            attempt_metric.application_blocking_time_ns += (  # type: ignore
+                                time.monotonic_ns() - block_time
+                            ) * 1000
                         break
                     c = await it.__anext__()
             except _ResetRow as e:

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

@@ -23,7 +23,6 @@
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT
 from google.cloud.bigtable.data.mutations import _EntryWithProto
 from google.cloud.bigtable.data._cross_sync import CrossSync
@@ -80,11 +79,10 @@ def __init__(
         self.is_retryable = retries.if_exception_type(
             *retryable_exceptions, bt_exceptions._MutateRowsIncomplete
         )
-        sleep_generator = TrackedBackoffGenerator(0.01, 2, 60)
         self._operation = lambda: CrossSync._Sync_Impl.retry_target(
             self._run_attempt,
             self.is_retryable,
-            sleep_generator,
+            metric.backoff_generator,
             operation_timeout,
             exception_factory=_retry_exception_factory,
         )
@@ -94,7 +92,6 @@ def __init__(
         self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries]
         self.remaining_indices = list(range(len(self.mutations)))
         self.errors: dict[int, list[Exception]] = {}
-        metric.backoff_generator = sleep_generator
         self._operation_metric = metric
 
     def start(self):

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

@@ -30,7 +30,6 @@
 from google.cloud.bigtable.data.exceptions import _ResetRow
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
-from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 from google.api_core import retry as retries
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
@@ -104,16 +103,14 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]:
 
         Yields:
             Row: The next row in the stream"""
-        self._operation_metric.backoff_generator = TrackedBackoffGenerator(
-            0.01, 60, multiplier=2
-        )
-        return CrossSync._Sync_Impl.retry_target_stream(
-            self._read_rows_attempt,
-            self._predicate,
-            self._operation_metric.backoff_generator,
-            self.operation_timeout,
-            exception_factory=_retry_exception_factory,
-        )
+        with self._operation_metric:
+            return CrossSync._Sync_Impl.retry_target_stream(
+                self._read_rows_attempt,
+                self._predicate,
+                self._operation_metric.backoff_generator,
+                self.operation_timeout,
+                exception_factory=_retry_exception_factory,
+            )
 
     def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]:
         """Attempt a single read_rows rpc call.
@@ -188,7 +185,7 @@ def chunk_stream(
     @staticmethod
     def merge_rows(
         chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None,
-        attempt_metric: ActiveAttemptMetric,
+        attempt_metric: ActiveAttemptMetric | None,
     ) -> CrossSync._Sync_Impl.Iterable[Row]:
         """Merge chunks into rows
 
@@ -199,7 +196,6 @@ def merge_rows(
         if chunks is None:
             return
         it = chunks.__iter__()
-        is_first_row = True
         while True:
             try:
                 c = it.__next__()
@@ -268,14 +264,12 @@ def merge_rows(
                         Cell(value, row_key, family, qualifier, ts, list(labels))
                     )
                     if c.commit_row:
-                        if is_first_row:
-                            is_first_row = False
-                            attempt_metric.attempt_first_response()
                         block_time = time.monotonic_ns()
                         yield Row(row_key, cells)
-                        attempt_metric.active_attempt.application_blocking_time_ns += (
-                            time.monotonic_ns() - block_time
-                        ) * 1000
+                        if attempt_metric is not None:
+                            attempt_metric.application_blocking_time_ns += (
+                                time.monotonic_ns() - block_time
+                            ) * 1000
                         break
                     c = it.__next__()
             except _ResetRow as e:

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

@@ -66,7 +66,6 @@
 from google.cloud.bigtable.data._helpers import _get_retryable_errors
 from google.cloud.bigtable.data._helpers import _get_timeouts
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry
 from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule
 from google.cloud.bigtable.data.row_filters import RowFilter
@@ -805,18 +804,17 @@ def read_rows_stream(
             operation_timeout, attempt_timeout, self
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
-        with self._metrics.create_operation(
-            OperationType.READ_ROWS, streaming=True
-        ) as operation_metric:
-            row_merger = CrossSync._Sync_Impl._ReadRowsOperation(
-                query,
-                self,
-                operation_timeout=operation_timeout,
-                attempt_timeout=attempt_timeout,
-                metric=operation_metric,
-                retryable_exceptions=retryable_excs,
-            )
-            return row_merger.start_operation()
+        row_merger = CrossSync._Sync_Impl._ReadRowsOperation(
+            query,
+            self,
+            operation_timeout=operation_timeout,
+            attempt_timeout=attempt_timeout,
+            metric=self._metrics.create_operation(
+                OperationType.READ_ROWS, streaming=True
+            ),
+            retryable_exceptions=retryable_excs,
+        )
+        return row_merger.start_operation()
 
     def read_rows(
         self,
@@ -906,22 +904,21 @@ def read_row(
             operation_timeout, attempt_timeout, self
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
-        with self._metrics.create_operation(
-            OperationType.READ_ROWS, streaming=False
-        ) as operation_metric:
-            row_merger = CrossSync._Sync_Impl._ReadRowsOperation(
-                query,
-                self,
-                operation_timeout=operation_timeout,
-                attempt_timeout=attempt_timeout,
-                metric=operation_metric,
-                retryable_exceptions=retryable_excs,
-            )
-            results_generator = row_merger.start_operation()
-            try:
-                return results_generator.__next__()
-            except StopIteration:
-                return None
+        row_merger = CrossSync._Sync_Impl._ReadRowsOperation(
+            query,
+            self,
+            operation_timeout=operation_timeout,
+            attempt_timeout=attempt_timeout,
+            metric=self._metrics.create_operation(
+                OperationType.READ_ROWS, streaming=False
+            ),
+            retryable_exceptions=retryable_excs,
+        )
+        results_generator = row_merger.start_operation()
+        try:
+            return results_generator.__next__()
+        except StopIteration:
+            return None
 
     def read_rows_sharded(
         self,
@@ -1101,10 +1098,9 @@ def sample_row_keys(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
         predicate = retries.if_exception_type(*retryable_excs)
-        sleep_generator = TrackedBackoffGenerator(0.01, 2, 60)
         with self._metrics.create_operation(
-            OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator
-        ):
+            OperationType.SAMPLE_ROW_KEYS
+        ) as operation_metric:
 
             def execute_rpc():
                 results = self.client._gapic_client.sample_row_keys(
@@ -1119,7 +1115,7 @@ def execute_rpc():
             return CrossSync._Sync_Impl.retry_target(
                 execute_rpc,
                 predicate,
-                sleep_generator,
+                operation_metric.backoff_generator,
                 operation_timeout,
                 exception_factory=_retry_exception_factory,
             )
@@ -1223,10 +1219,9 @@ def mutate_row(
             )
         else:
             predicate = retries.if_exception_type()
-        sleep_generator = TrackedBackoffGenerator(0.01, 2, 60)
         with self._metrics.create_operation(
-            OperationType.MUTATE_ROW, backoff_generator=sleep_generator
-        ):
+            OperationType.MUTATE_ROW
+        ) as operation_metric:
             target = partial(
                 self.client._gapic_client.mutate_row,
                 request=MutateRowRequest(
@@ -1243,7 +1238,7 @@ def mutate_row(
             return CrossSync._Sync_Impl.retry_target(
                 target,
                 predicate,
-                sleep_generator,
+                operation_metric.backoff_generator,
                 operation_timeout,
                 exception_factory=_retry_exception_factory,
             )

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

@@ -15,6 +15,7 @@
 # This file is automatically generated by CrossSync. Do not edit manually.
 
 from __future__ import annotations
+import time
 from functools import wraps
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
@@ -76,7 +77,8 @@ def register_operation(self, operation):
         operation.handlers.append(self)
 
     def on_operation_complete(self, op):
-        del self.operation_map[op.uuid]
+        if op.uuid in self.operation_map:
+            del self.operation_map[op.uuid]
 
     def on_operation_cancelled(self, op):
         self.on_operation_complete(op)
@@ -105,9 +107,15 @@ 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
             encountered_exc = None
             try:
                 for response in call:
+                    if not has_first_response:
+                        operation.first_response_latency_ns = (
+                            time.monotonic_ns() - operation.start_time_ns
+                        )
+                        has_first_response = True
                     yield response
             except Exception as e:
                 encountered_exc = e