followed same operation management pattern for read_rows as mutate_rows

Author: daniel-sanche

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

@@ -118,13 +118,14 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
         Yields:
             Row: The next row in the stream
         """
-        return CrossSync.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.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.Iterable[Row]:
         """

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

@@ -1020,18 +1020,17 @@ async def read_rows_stream(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
 
-        with self._metrics.create_operation(
-            OperationType.READ_ROWS, streaming=True
-        ) as operation_metric:
-            row_merger = CrossSync._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._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()
 
     @CrossSync.convert
     async def read_rows(
@@ -1129,22 +1128,21 @@ async def read_row(
         )
         retryable_excs = _get_retryable_errors(retryable_errors, self)
 
-        with self._metrics.create_operation(
-            OperationType.READ_ROWS, streaming=False
-        ) as operation_metric:
-            row_merger = CrossSync._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.__anext__()
-            except StopAsyncIteration:
-                return None
+        row_merger = CrossSync._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.__anext__()
+        except StopAsyncIteration:
+            return None
 
     @CrossSync.convert
     async def read_rows_sharded(