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

Commit 14d252b

Browse files
daniel-sanchedaniel-sanche
committed
fixed operation end for read_rows
1 parent 2adec5a commit 14d252b

2 files changed

Lines changed: 245 additions & 227 deletions

File tree

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

Lines changed: 131 additions & 121 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from __future__ import annotations
1717

18-
from typing import Sequence, TYPE_CHECKING
18+
from typing import Callable, Sequence, TYPE_CHECKING
1919

2020
import time
2121

@@ -117,14 +117,13 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
117117
Yields:
118118
Row: The next row in the stream
119119
"""
120-
with self._operation_metric:
121-
return CrossSync.retry_target_stream(
122-
self._read_rows_attempt,
123-
self._predicate,
124-
self._operation_metric.backoff_generator,
125-
self.operation_timeout,
126-
exception_factory=_retry_exception_factory,
127-
)
120+
return CrossSync.retry_target_stream(
121+
self._read_rows_attempt,
122+
self._predicate,
123+
self._operation_metric.backoff_generator,
124+
self.operation_timeout,
125+
exception_factory=_retry_exception_factory,
126+
)
128127

129128
def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
130129
"""
@@ -136,7 +135,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
136135
Yields:
137136
Row: The next row in the stream
138137
"""
139-
attempt_metric = self._operation_metric.start_attempt()
138+
self._operation_metric.start_attempt()
140139
# revise request keys and ranges between attempts
141140
if self._last_yielded_row_key is not None:
142141
# if this is a retry, try to trim down the request to avoid ones we've already processed
@@ -147,20 +146,20 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
147146
)
148147
except _RowSetComplete:
149148
# if we've already seen all the rows, we're done
150-
return self.merge_rows(None, attempt_metric)
149+
return self.merge_rows(None, self._operation_metric, self._predicate)
151150
# revise the limit based on number of rows already yielded
152151
if self._remaining_count is not None:
153152
self.request.rows_limit = self._remaining_count
154153
if self._remaining_count == 0:
155-
return self.merge_rows(None, attempt_metric)
154+
return self.merge_rows(None, self._operation_metric, self._predicate)
156155
# create and return a new row merger
157156
gapic_stream = self.target.client._gapic_client.read_rows(
158157
self.request,
159158
timeout=next(self.attempt_timeout_gen),
160159
retry=None,
161160
)
162161
chunked_stream = self.chunk_stream(gapic_stream)
163-
return self.merge_rows(chunked_stream, attempt_metric)
162+
return self.merge_rows(chunked_stream, self._operation_metric, self._predicate)
164163

165164
@CrossSync.convert()
166165
async def chunk_stream(
@@ -220,7 +219,8 @@ async def chunk_stream(
220219
)
221220
async def merge_rows(
222221
chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None,
223-
attempt_metric: ActiveAttemptMetric | None,
222+
operation_metric: ActiveOperationMetric,
223+
retryable_predicate: Callable[[Exception], bool],
224224
) -> CrossSync.Iterable[Row]:
225225
"""
226226
Merge chunks into rows
@@ -230,115 +230,125 @@ async def merge_rows(
230230
Yields:
231231
Row: the next row in the stream
232232
"""
233-
if chunks is None:
234-
return
235-
it = chunks.__aiter__()
236-
# For each row
237-
while True:
238-
try:
239-
c = await it.__anext__()
240-
except CrossSync.StopIteration:
241-
# stream complete
233+
try:
234+
if chunks is None:
235+
operation_metric.end_with_success()
242236
return
243-
row_key = c.row_key
244-
245-
if not row_key:
246-
raise InvalidChunk("first row chunk is missing key")
247-
248-
cells = []
249-
250-
# shared per cell storage
251-
family: str | None = None
252-
qualifier: bytes | None = None
253-
254-
try:
255-
# for each cell
256-
while True:
257-
if c.reset_row:
258-
raise _ResetRow(c)
259-
k = c.row_key
260-
f = c.family_name.value
261-
q = c.qualifier.value if c.HasField("qualifier") else None
262-
if k and k != row_key:
263-
raise InvalidChunk("unexpected new row key")
264-
if f:
265-
family = f
266-
if q is not None:
267-
qualifier = q
268-
else:
269-
raise InvalidChunk("new family without qualifier")
270-
elif family is None:
271-
raise InvalidChunk("missing family")
272-
elif q is not None:
273-
if family is None:
274-
raise InvalidChunk("new qualifier without family")
275-
qualifier = q
276-
elif qualifier is None:
277-
raise InvalidChunk("missing qualifier")
278-
279-
ts = c.timestamp_micros
280-
labels = c.labels if c.labels else []
281-
value = c.value
282-
283-
# merge split cells
284-
if c.value_size > 0:
285-
buffer = [value]
286-
while c.value_size > 0:
287-
# throws when premature end
288-
c = await it.__anext__()
289-
290-
t = c.timestamp_micros
291-
cl = c.labels
292-
k = c.row_key
293-
if (
294-
c.HasField("family_name")
295-
and c.family_name.value != family
296-
):
297-
raise InvalidChunk("family changed mid cell")
298-
if (
299-
c.HasField("qualifier")
300-
and c.qualifier.value != qualifier
301-
):
302-
raise InvalidChunk("qualifier changed mid cell")
303-
if t and t != ts:
304-
raise InvalidChunk("timestamp changed mid cell")
305-
if cl and cl != labels:
306-
raise InvalidChunk("labels changed mid cell")
307-
if k and k != row_key:
308-
raise InvalidChunk("row key changed mid cell")
309-
310-
if c.reset_row:
311-
raise _ResetRow(c)
312-
buffer.append(c.value)
313-
value = b"".join(buffer)
314-
cells.append(
315-
Cell(value, row_key, family, qualifier, ts, list(labels))
316-
)
317-
if c.commit_row:
318-
block_time = time.monotonic_ns()
319-
yield Row(row_key, cells)
320-
# most metric operations use setters, but this one updates
321-
# the value directly to avoid extra overhead
322-
if attempt_metric is not None:
323-
attempt_metric.application_blocking_time_ns += ( # type: ignore
324-
time.monotonic_ns() - block_time
325-
) * 1000
326-
break
237+
it = chunks.__aiter__()
238+
# For each row
239+
while True:
240+
try:
327241
c = await it.__anext__()
328-
except _ResetRow as e:
329-
c = e.chunk
330-
if (
331-
c.row_key
332-
or c.HasField("family_name")
333-
or c.HasField("qualifier")
334-
or c.timestamp_micros
335-
or c.labels
336-
or c.value
337-
):
338-
raise InvalidChunk("reset row with data")
339-
continue
340-
except CrossSync.StopIteration:
341-
raise InvalidChunk("premature end of stream")
242+
except CrossSync.StopIteration:
243+
# stream complete
244+
operation_metric.end_with_success()
245+
return
246+
row_key = c.row_key
247+
248+
if not row_key:
249+
raise InvalidChunk("first row chunk is missing key")
250+
251+
cells = []
252+
253+
# shared per cell storage
254+
family: str | None = None
255+
qualifier: bytes | None = None
256+
257+
try:
258+
# for each cell
259+
while True:
260+
if c.reset_row:
261+
raise _ResetRow(c)
262+
k = c.row_key
263+
f = c.family_name.value
264+
q = c.qualifier.value if c.HasField("qualifier") else None
265+
if k and k != row_key:
266+
raise InvalidChunk("unexpected new row key")
267+
if f:
268+
family = f
269+
if q is not None:
270+
qualifier = q
271+
else:
272+
raise InvalidChunk("new family without qualifier")
273+
elif family is None:
274+
raise InvalidChunk("missing family")
275+
elif q is not None:
276+
if family is None:
277+
raise InvalidChunk("new qualifier without family")
278+
qualifier = q
279+
elif qualifier is None:
280+
raise InvalidChunk("missing qualifier")
281+
282+
ts = c.timestamp_micros
283+
labels = c.labels if c.labels else []
284+
value = c.value
285+
286+
# merge split cells
287+
if c.value_size > 0:
288+
buffer = [value]
289+
while c.value_size > 0:
290+
# throws when premature end
291+
c = await it.__anext__()
292+
293+
t = c.timestamp_micros
294+
cl = c.labels
295+
k = c.row_key
296+
if (
297+
c.HasField("family_name")
298+
and c.family_name.value != family
299+
):
300+
raise InvalidChunk("family changed mid cell")
301+
if (
302+
c.HasField("qualifier")
303+
and c.qualifier.value != qualifier
304+
):
305+
raise InvalidChunk("qualifier changed mid cell")
306+
if t and t != ts:
307+
raise InvalidChunk("timestamp changed mid cell")
308+
if cl and cl != labels:
309+
raise InvalidChunk("labels changed mid cell")
310+
if k and k != row_key:
311+
raise InvalidChunk("row key changed mid cell")
312+
313+
if c.reset_row:
314+
raise _ResetRow(c)
315+
buffer.append(c.value)
316+
value = b"".join(buffer)
317+
cells.append(
318+
Cell(value, row_key, family, qualifier, ts, list(labels))
319+
)
320+
if c.commit_row:
321+
block_time = time.monotonic_ns()
322+
yield Row(row_key, cells)
323+
# most metric operations use setters, but this one updates
324+
# the value directly to avoid extra overhead
325+
if operation_metric.active_attempt is not None:
326+
operation_metric.active_attempt.application_blocking_time_ns += ( # type: ignore
327+
time.monotonic_ns() - block_time
328+
) * 1000
329+
break
330+
c = await it.__anext__()
331+
except _ResetRow as e:
332+
c = e.chunk
333+
if (
334+
c.row_key
335+
or c.HasField("family_name")
336+
or c.HasField("qualifier")
337+
or c.timestamp_micros
338+
or c.labels
339+
or c.value
340+
):
341+
raise InvalidChunk("reset row with data")
342+
continue
343+
except CrossSync.StopIteration:
344+
raise InvalidChunk("premature end of stream")
345+
except Exception as generic_exception:
346+
if not retryable_predicate(generic_exception):
347+
operation_metric.end_attempt_with_status(generic_exception)
348+
raise generic_exception
349+
else:
350+
operation_metric.end_with_success()
351+
342352

343353
@staticmethod
344354
def _revise_request_rowset(

0 commit comments

Comments
 (0)