fixed lint and mypy

Author: daniel-sanche

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

@@ -169,7 +169,7 @@ async def _run_attempt(self):
                 retry after the attempt is complete
             GoogleAPICallError: if the gapic rpc fails
         """
-         # register attempt start
+        # register attempt start
         self._operation_metric.start_attempt()
         request_entries = [self.mutations[idx].proto for idx in self.remaining_indices]
         # track mutations in this request that have not been finalized yet

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

@@ -34,11 +34,11 @@
 from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
 
 from google.api_core import retry as retries
-from google.api_core.retry import exponential_sleep_generator
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 
 if TYPE_CHECKING:
+    from google.cloud.bigtable.data._metrics import ActiveAttemptMetric
     from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 
     if CrossSync.is_async:
@@ -139,7 +139,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
         Yields:
             Row: The next row in the stream
         """
-        self._operation_metric.start_attempt()
+        attempt_metric = self._operation_metric.start_attempt()
         # revise request keys and ranges between attempts
         if self._last_yielded_row_key is not None:
             # if this is a retry, try to trim down the request to avoid ones we've already processed
@@ -150,20 +150,20 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
                 )
             except _RowSetComplete:
                 # if we've already seen all the rows, we're done
-                return self.merge_rows(None, self._operation_metric)
+                return self.merge_rows(None, attempt_metric)
         # revise the limit based on number of rows already yielded
         if self._remaining_count is not None:
             self.request.rows_limit = self._remaining_count
             if self._remaining_count == 0:
-                return self.merge_rows(None, self._operation_metric)
+                return self.merge_rows(None, attempt_metric)
         # create and return a new row merger
         gapic_stream = self.target.client._gapic_client.read_rows(
             self.request,
             timeout=next(self.attempt_timeout_gen),
             retry=None,
         )
         chunked_stream = self.chunk_stream(gapic_stream)
-        return self.merge_rows(chunked_stream, self._operation_metric)
+        return self.merge_rows(chunked_stream, attempt_metric)
 
     @CrossSync.convert()
     async def chunk_stream(
@@ -223,7 +223,7 @@ async def chunk_stream(
     )
     async def merge_rows(
         chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None,
-        operation_metric: ActiveOperationMetric,
+        attempt_metric: ActiveAttemptMetric,
     ) -> CrossSync.Iterable[Row]:
         """
         Merge chunks into rows
@@ -322,13 +322,13 @@ async def merge_rows(
                         if is_first_row:
                             # record first row latency in metrics
                             is_first_row = False
-                            operation_metric.attempt_first_response()
-                        block_time = time.monotonic()
+                            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
-                        operation_metric.active_attempt.application_blocking_time_ms += (  # type: ignore
-                            time.monotonic() - block_time
+                        attempt_metric.active_attempt.application_blocking_time_ns += (  # type: ignore
+                            time.monotonic_ns() - block_time
                         ) * 1000
                         break
                     c = await it.__anext__()

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

@@ -125,7 +125,6 @@
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
     from google.cloud.bigtable.data._helpers import ShardedQuery
-    from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 
     if CrossSync.is_async:
         from google.cloud.bigtable.data._async.mutations_batcher import (
@@ -1021,7 +1020,7 @@ async def read_rows_stream(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
 
-        with self._metrics.create_operation(
+        async with self._metrics.create_operation(
             OperationType.READ_ROWS, streaming=True
         ) as operation_metric:
             row_merger = CrossSync._ReadRowsOperation(
@@ -1130,7 +1129,7 @@ async def read_row(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
 
-        with self._metrics.create_operation(
+        async with self._metrics.create_operation(
             OperationType.READ_ROWS, streaming=False
         ) as operation_metric:
             row_merger = CrossSync._ReadRowsOperation(

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

@@ -182,19 +182,21 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]
             yield mutations[start_idx:end_idx]
 
     async def add_to_flow_with_metrics(
-        self, mutations: RowMutationEntry | list[RowMutationEntry], metrics_controller: BigtableClientSideMetricsController
+        self,
+        mutations: RowMutationEntry | list[RowMutationEntry],
+        metrics_controller: BigtableClientSideMetricsController,
     ):
         inner_generator = self.add_to_flow(mutations)
         while True:
             # start a new metric
             metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS)
-            flow_start_time = time.monotonic()
+            flow_start_time = time.monotonic_ns()
             try:
                 value = await inner_generator.__anext__()
             except StopAsyncIteration:
                 metric.cancel()
                 return
-            metric.flow_throttling_time = time.monotonic() - flow_start_time
+            metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time
             yield value, metric
 
 
@@ -373,9 +375,14 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
         """
         # flush new entries
         in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = []
-        async for batch, metric in self._flow_control.add_to_flow_with_metrics(new_entries, self._target._metrics):
+        async for batch, metric in self._flow_control.add_to_flow_with_metrics(
+            new_entries, self._target._metrics
+        ):
             batch_task = CrossSync.create_task(
-                self._execute_mutate_rows, batch, metric, sync_executor=self._sync_rpc_executor
+                self._execute_mutate_rows,
+                batch,
+                metric,
+                sync_executor=self._sync_rpc_executor,
             )
             in_process_requests.append(batch_task)
         # wait for all inflight requests to complete

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

@@ -13,7 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
-from typing import Callable, Tuple, cast, TYPE_CHECKING
+from typing import Tuple, cast, TYPE_CHECKING
 
 import time
 import re
@@ -184,7 +184,7 @@ def start(self) -> None:
             return self._handle_error(INVALID_STATE_ERROR.format("start", self.state))
         self.start_time_ns = time.monotonic_ns()
 
-    def start_attempt(self) -> None:
+    def start_attempt(self) -> ActiveAttemptMetric:
         """
         Called to initiate a new attempt for the operation.
 
@@ -210,6 +210,7 @@ def start_attempt(self) -> None:
             backoff_ns = 0
 
         self.active_attempt = ActiveAttemptMetric(backoff_before_attempt_ns=backoff_ns)
+        return self.active_attempt
 
     def add_response_metadata(self, metadata: dict[str, bytes | str]) -> None:
         """
@@ -428,4 +429,4 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
         if exc_val is None:
             self.end_with_success()
         else:
-            self.end_with_status(exc_val)
+            self.end_with_status(exc_val)

google/cloud/bigtable/data/_metrics/handlers/_base.py

@@ -35,4 +35,4 @@ def on_operation_cancelled(self, op: ActiveOperationMetric) -> None:
     def on_attempt_complete(
         self, attempt: CompletedAttemptMetric, op: ActiveOperationMetric
     ) -> None:
-        pass
+        pass

google/cloud/bigtable/data/_metrics/handlers/_stdout.py

@@ -45,4 +45,4 @@ def print(self):
             print(
                 f"{ops_type}: count: {count}, avg latency: {avg_latency:.2f} ms, avg attempts: {avg_attempts:.1f}"
             )
-        print()
+        print()

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

@@ -23,12 +23,14 @@
 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
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data.mutations import RowMutationEntry
+    from google.cloud.bigtable.data._metrics import ActiveOperationMetric
     from google.cloud.bigtable_v2.services.bigtable.client import (
         BigtableClient as GapicClientType,
     )
@@ -54,6 +56,8 @@ class _MutateRowsOperation:
         operation_timeout: the timeout to use for the entire operation, in seconds.
         attempt_timeout: the timeout to use for each mutate_rows attempt, in seconds.
             If not specified, the request will run until operation_timeout is reached.
+        metric: the metric object representing the active operation
+        retryable_exceptions: a list of exceptions that should be retried
     """
 
     def __init__(
@@ -63,6 +67,7 @@ def __init__(
         mutation_entries: list["RowMutationEntry"],
         operation_timeout: float,
         attempt_timeout: float | None,
+        metric: ActiveOperationMetric,
         retryable_exceptions: Sequence[type[Exception]] = (),
     ):
         total_mutations = sum((len(entry.mutations) for entry in mutation_entries))
@@ -75,7 +80,7 @@ def __init__(
         self.is_retryable = retries.if_exception_type(
             *retryable_exceptions, bt_exceptions._MutateRowsIncomplete
         )
-        sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60)
+        sleep_generator = TrackedBackoffGenerator(0.01, 2, 60)
         self._operation = lambda: CrossSync._Sync_Impl.retry_target(
             self._run_attempt,
             self.is_retryable,
@@ -89,37 +94,40 @@ 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):
         """Start the operation, and run until completion
 
         Raises:
             MutationsExceptionGroup: if any mutations failed"""
-        try:
-            self._operation()
-        except Exception as exc:
-            incomplete_indices = self.remaining_indices.copy()
-            for idx in incomplete_indices:
-                self._handle_entry_error(idx, exc)
-        finally:
-            all_errors: list[Exception] = []
-            for idx, exc_list in self.errors.items():
-                if len(exc_list) == 0:
-                    raise core_exceptions.ClientError(
-                        f"Mutation {idx} failed with no associated errors"
+        with self._operation_metric:
+            try:
+                self._operation()
+            except Exception as exc:
+                incomplete_indices = self.remaining_indices.copy()
+                for idx in incomplete_indices:
+                    self._handle_entry_error(idx, exc)
+            finally:
+                all_errors: list[Exception] = []
+                for idx, exc_list in self.errors.items():
+                    if len(exc_list) == 0:
+                        raise core_exceptions.ClientError(
+                            f"Mutation {idx} failed with no associated errors"
+                        )
+                    elif len(exc_list) == 1:
+                        cause_exc = exc_list[0]
+                    else:
+                        cause_exc = bt_exceptions.RetryExceptionGroup(exc_list)
+                    entry = self.mutations[idx].entry
+                    all_errors.append(
+                        bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc)
+                    )
+                if all_errors:
+                    raise bt_exceptions.MutationsExceptionGroup(
+                        all_errors, len(self.mutations)
                     )
-                elif len(exc_list) == 1:
-                    cause_exc = exc_list[0]
-                else:
-                    cause_exc = bt_exceptions.RetryExceptionGroup(exc_list)
-                entry = self.mutations[idx].entry
-                all_errors.append(
-                    bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc)
-                )
-            if all_errors:
-                raise bt_exceptions.MutationsExceptionGroup(
-                    all_errors, len(self.mutations)
-                )
 
     def _run_attempt(self):
         """Run a single attempt of the mutate_rows rpc.
@@ -128,6 +136,7 @@ def _run_attempt(self):
             _MutateRowsIncomplete: if there are failed mutations eligible for
                 retry after the attempt is complete
             GoogleAPICallError: if the gapic rpc fails"""
+        self._operation_metric.start_attempt()
         request_entries = [self.mutations[idx].proto for idx in self.remaining_indices]
         active_request_indices = {
             req_idx: orig_idx