Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 88644b2

Browse files
daniel-sanchedaniel-sanche
committed
fixed lint
1 parent 2e0e402 commit 88644b2

13 files changed

Lines changed: 180 additions & 120 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,9 @@ def __init__(
106106
self.is_retryable,
107107
metric.backoff_generator,
108108
operation_timeout,
109-
exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
109+
exception_factory=self._operation_metric.track_terminal_error(
110+
_retry_exception_factory
111+
),
110112
on_error=self._operation_metric.track_retryable_error,
111113
)
112114
# initialize state

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,9 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
123123
self._predicate,
124124
self._operation_metric.backoff_generator,
125125
self.operation_timeout,
126-
exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory),
126+
exception_factory=self._operation_metric.track_terminal_error(
127+
_retry_exception_factory
128+
),
127129
on_error=self._operation_metric.track_retryable_error,
128130
)
129131

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1389,7 +1389,9 @@ async def execute_rpc():
13891389
predicate,
13901390
operation_metric.backoff_generator,
13911391
operation_timeout,
1392-
exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
1392+
exception_factory=operation_metric.track_terminal_error(
1393+
_retry_exception_factory
1394+
),
13931395
on_error=operation_metric.track_retryable_error,
13941396
)
13951397

@@ -1523,7 +1525,9 @@ async def mutate_row(
15231525
predicate,
15241526
operation_metric.backoff_generator,
15251527
operation_timeout,
1526-
exception_factory=operation_metric.track_terminal_error(_retry_exception_factory),
1528+
exception_factory=operation_metric.track_terminal_error(
1529+
_retry_exception_factory
1530+
),
15271531
on_error=operation_metric.track_retryable_error,
15281532
)
15291533

@@ -1638,7 +1642,7 @@ async def check_and_mutate_row(
16381642
false_case_mutations = [false_case_mutations]
16391643
false_case_list = [m._to_pb() for m in false_case_mutations or []]
16401644

1641-
with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op:
1645+
with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
16421646
result = await self.client._gapic_client.check_and_mutate_row(
16431647
request=CheckAndMutateRowRequest(
16441648
true_mutations=true_case_list,
@@ -1696,7 +1700,7 @@ async def read_modify_write_row(
16961700
if not rules:
16971701
raise ValueError("rules must contain at least one item")
16981702

1699-
with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op:
1703+
with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
17001704
result = await self.client._gapic_client.read_modify_write_row(
17011705
request=ReadModifyWriteRowRequest(
17021706
rules=[rule._to_pb() for rule in rules],

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

Lines changed: 58 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717

1818
import time
1919
from functools import wraps
20-
from grpc import StatusCode
2120

2221
from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
2322
from google.cloud.bigtable.data._metrics.data_model import OperationState
@@ -63,21 +62,28 @@ def wrapper(self, continuation, client_call_details, request):
6362

6463
return wrapper
6564

65+
6666
@CrossSync.convert
67-
async def _get_metadata(source) -> dict[str, str|bytes] | None:
67+
async def _get_metadata(source) -> dict[str, str | bytes] | None:
6868
"""Helper to extract metadata from a call or RpcError"""
6969
try:
7070
if CrossSync.is_async:
7171
# grpc.aio returns metadata in Metadata objects
7272
if isinstance(source, AioRpcError):
73-
metadata = list(source.trailing_metadata()) + list(source.initial_metadata())
73+
metadata = list(source.trailing_metadata()) + list(
74+
source.initial_metadata()
75+
)
7476
else:
75-
metadata = list(await source.trailing_metadata()) + list(await source.initial_metadata())
77+
metadata = list(await source.trailing_metadata()) + list(
78+
await source.initial_metadata()
79+
)
7680
else:
7781
# sync grpc returns metadata as a sequence of tuples
78-
metadata: Sequence[tuple[str. str|bytes]] = source.trailing_metadata() + source.initial_metadata()
82+
metadata: Sequence[tuple[str.str | bytes]] = (
83+
source.trailing_metadata() + source.initial_metadata()
84+
)
7985
# convert metadata to dict format
80-
return {k:v for k,v in metadata}
86+
return {k: v for k, v in metadata}
8187
except Exception:
8288
# ignore errors while fetching metadata
8389
return None
@@ -96,6 +102,12 @@ class AsyncBigtableMetricsInterceptor(
96102
async def intercept_unary_unary(
97103
self, operation, continuation, client_call_details, request
98104
):
105+
"""
106+
Interceptor for unary rpcs:
107+
- MutateRow
108+
- CheckAndMutateRow
109+
- ReadModifyWriteRow
110+
"""
99111
metadata = None
100112
try:
101113
call = await continuation(client_call_details, request)
@@ -113,36 +125,49 @@ async def intercept_unary_unary(
113125
async def intercept_unary_stream(
114126
self, operation, continuation, client_call_details, request
115127
):
116-
async def response_wrapper(call):
117-
# only track has_first response for READ_ROWS
118-
has_first_response = (
119-
operation.first_response_latency_ns is not None
120-
or operation.op_type != OperationType.READ_ROWS
121-
)
122-
encountered_exc = None
123-
try:
124-
async for response in call:
125-
# record time to first response. Currently only used for READ_ROWs
126-
if not has_first_response:
127-
operation.first_response_latency_ns = (
128-
time.monotonic_ns() - operation.start_time_ns
129-
)
130-
has_first_response = True
131-
yield response
132-
except Exception as e:
133-
# handle errors while processing stream
134-
encountered_exc = e
135-
raise
136-
finally:
137-
if call is not None:
138-
metadata = await _get_metadata(encountered_exc or call)
139-
if metadata is not None:
140-
operation.add_response_metadata(metadata)
141-
128+
"""
129+
Interceptor for streaming rpcs:
130+
- ReadRows
131+
- MutateRows
132+
- SampleRowKeys
133+
"""
142134
try:
143-
return response_wrapper(await continuation(client_call_details, request))
135+
return self._streaming_generator_wrapper(
136+
operation, await continuation(client_call_details, request)
137+
)
144138
except Exception as rpc_error:
145139
metadata = await _get_metadata(rpc_error)
146140
if metadata is not None:
147141
operation.add_response_metadata(metadata)
148142
raise rpc_error
143+
144+
@staticmethod
145+
@CrossSync.convert
146+
async def _streaming_generator_wrapper(operation, call):
147+
"""
148+
Wrapped generator to be returned by intercept_unary_stream
149+
"""
150+
# only track has_first response for READ_ROWS
151+
has_first_response = (
152+
operation.first_response_latency_ns is not None
153+
or operation.op_type != OperationType.READ_ROWS
154+
)
155+
encountered_exc = None
156+
try:
157+
async for response in call:
158+
# record time to first response. Currently only used for READ_ROWs
159+
if not has_first_response:
160+
operation.first_response_latency_ns = (
161+
time.monotonic_ns() - operation.start_time_ns
162+
)
163+
has_first_response = True
164+
yield response
165+
except Exception as e:
166+
# handle errors while processing stream
167+
encountered_exc = e
168+
raise
169+
finally:
170+
if call is not None:
171+
metadata = await _get_metadata(encountered_exc or call)
172+
if metadata is not None:
173+
operation.add_response_metadata(metadata)

‎google/cloud/bigtable/data/_helpers.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
"""
1717
from __future__ import annotations
1818

19-
from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union
19+
from typing import Sequence, List, Tuple, TYPE_CHECKING, Union
2020
import time
2121
import enum
2222
from collections import namedtuple

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

Lines changed: 22 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,6 @@
1515

1616
from typing import ClassVar, Tuple, cast, TYPE_CHECKING
1717

18-
1918
import time
2019
import re
2120
import logging
@@ -36,7 +35,6 @@
3635
from google.cloud.bigtable_v2.types.response_params import ResponseParams
3736
from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator
3837
from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete
39-
from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
4038
from google.protobuf.message import DecodeError
4139

4240
if TYPE_CHECKING:
@@ -169,8 +167,9 @@ class ActiveOperationMetric:
169167
# time waiting on flow control, in nanoseconds
170168
flow_throttling_time_ns: int = 0
171169

172-
173-
_active_operation_context: ClassVar[contextvars.ContextVar] = contextvars.ContextVar("active_operation_context")
170+
_active_operation_context: ClassVar[
171+
contextvars.ContextVar
172+
] = contextvars.ContextVar("active_operation_context")
174173

175174
@classmethod
176175
def get_active(cls):
@@ -403,10 +402,7 @@ def track_retryable_error(self, exc: Exception) -> None:
403402
"""
404403
try:
405404
# record metadata from failed rpc
406-
if (
407-
isinstance(exc, GoogleAPICallError)
408-
and exc.errors
409-
):
405+
if isinstance(exc, GoogleAPICallError) and exc.errors:
410406
rpc_error = exc.errors[-1]
411407
metadata = list(rpc_error.trailing_metadata()) + list(
412408
rpc_error.initial_metadata()
@@ -422,24 +418,28 @@ def track_retryable_error(self, exc: Exception) -> None:
422418
else:
423419
self.end_attempt_with_status(exc)
424420

425-
def track_terminal_error(self, exception_factory:callable[
426-
[list[Exception], RetryFailureReason, float | None],tuple[Exception, Exception | None],
427-
]) -> callable[[list[Exception], RetryFailureReason, float | None], None]:
421+
def track_terminal_error(
422+
self,
423+
exception_factory: callable[
424+
[list[Exception], RetryFailureReason, float | None],
425+
tuple[Exception, Exception | None],
426+
],
427+
) -> callable[[list[Exception], RetryFailureReason, float | None], None]:
428428
"""
429429
Used as input to api_core.Retry classes, to track when terminal errors are encountered
430430
431431
Should be used as a wrapper over an exception_factory callback
432432
"""
433+
433434
def wrapper(
434-
exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None
435+
exc_list: list[Exception],
436+
reason: RetryFailureReason,
437+
timeout_val: float | None,
435438
) -> tuple[Exception, Exception | None]:
436439
source_exc, cause_exc = exception_factory(exc_list, reason, timeout_val)
437440
try:
438441
# record metadata from failed rpc
439-
if (
440-
isinstance(source_exc, GoogleAPICallError)
441-
and source_exc.errors
442-
):
442+
if isinstance(source_exc, GoogleAPICallError) and source_exc.errors:
443443
rpc_error = source_exc.errors[-1]
444444
metadata = list(rpc_error.trailing_metadata()) + list(
445445
rpc_error.initial_metadata()
@@ -448,12 +448,17 @@ def wrapper(
448448
except Exception:
449449
# ignore errors in metadata collection
450450
pass
451-
if reason == RetryFailureReason.TIMEOUT and self.state == OperationState.ACTIVE_ATTEMPT and exc_list:
451+
if (
452+
reason == RetryFailureReason.TIMEOUT
453+
and self.state == OperationState.ACTIVE_ATTEMPT
454+
and exc_list
455+
):
452456
# record ending attempt for timeout failures
453457
attempt_exc = exc_list[-1]
454458
self.track_retryable_error(attempt_exc)
455459
self.end_with_status(source_exc)
456460
return source_exc, cause_exc
461+
457462
return wrapper
458463

459464
@staticmethod

‎google/cloud/bigtable/data/_metrics/metrics_controller.py‎

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,20 +13,10 @@
1313
# limitations under the License.
1414
from __future__ import annotations
1515

16-
from typing import TYPE_CHECKING
17-
1816
from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
1917
from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
2018
from google.cloud.bigtable.data._metrics.data_model import OperationType
2119

22-
if TYPE_CHECKING:
23-
from google.cloud.bigtable.data._async.metrics_interceptor import (
24-
AsyncBigtableMetricsInterceptor,
25-
)
26-
from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import (
27-
BigtableMetricsInterceptor,
28-
)
29-
3020

3121
class BigtableClientSideMetricsController:
3222
"""

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1377,7 +1377,7 @@ def check_and_mutate_row(
13771377
):
13781378
false_case_mutations = [false_case_mutations]
13791379
false_case_list = [m._to_pb() for m in false_case_mutations or []]
1380-
with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op:
1380+
with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
13811381
result = self.client._gapic_client.check_and_mutate_row(
13821382
request=CheckAndMutateRowRequest(
13831383
true_mutations=true_case_list,
@@ -1431,7 +1431,7 @@ def read_modify_write_row(
14311431
rules = [rules]
14321432
if not rules:
14331433
raise ValueError("rules must contain at least one item")
1434-
with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op:
1434+
with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
14351435
result = self.client._gapic_client.read_modify_write_row(
14361436
request=ReadModifyWriteRowRequest(
14371437
rules=[rule._to_pb() for rule in rules],

0 commit comments

Comments
 (0)