Skip to content

fix(bigtable): preserve row key across ReadRowsResponse boundaries in chunk_stream #18748

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 1 commit into
base: main
Choose a base branch
from
+191 −334
Open

fix(bigtable): preserve row key across ReadRowsResponse boundaries in chunk_stream #18748

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,7 @@ async def chunk_stream(
Yields:
ReadRowsResponsePB.CellChunk: the next chunk in the stream
"""
current_key = None
async for resp in await stream:
# extract proto from proto-plus wrapper
resp = resp._pb
Expand All @@ -190,12 +191,11 @@ async def chunk_stream(
raise InvalidChunk("last scanned out of order")
self._last_yielded_row_key = resp.last_scanned_row_key

current_key = None
# process each chunk in the response
for c in resp.chunks:
if current_key is None:
current_key = c.row_key
if current_key is None:
if not current_key:
raise InvalidChunk("first chunk is missing a row key")
elif (
self._last_yielded_row_key
Expand Down
207 changes: 64 additions & 143 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,36 +17,22 @@
# This file is automatically generated by CrossSync. Do not edit manually.

from __future__ import annotations

import time
from typing import TYPE_CHECKING, Sequence

from grpc import StatusCode

from google.cloud.bigtable.data._cross_sync import CrossSync
from google.cloud.bigtable.data._helpers import (
_attempt_timeout_generator,
_rst_stream_aware_predicate,
)
from google.cloud.bigtable.data._helpers import _attempt_timeout_generator, _rst_stream_aware_predicate
from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
from google.cloud.bigtable.data.exceptions import (
InvalidChunk,
_ResetRow,
_RowSetComplete,
)
from google.cloud.bigtable.data.exceptions import InvalidChunk, _ResetRow, _RowSetComplete
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
from google.cloud.bigtable.data.row import Cell, Row
from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB
from google.cloud.bigtable_v2.types import RowRange as RowRangePB
from google.cloud.bigtable_v2.types import RowSet as RowSetPB

if TYPE_CHECKING:
from google.cloud.bigtable.data._metrics import ActiveOperationMetric
from google.cloud.bigtable.data._sync_autogen.client import (
_DataApiTarget as TargetType,
)

from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget as TargetType

class _ReadRowsOperation:
"""
Expand All @@ -68,35 +54,13 @@ class _ReadRowsOperation:
metric: the metric object representing the active operation
retryable_exceptions: A list of exceptions that should trigger a retry
"""
__slots__ = ('attempt_timeout_gen', 'operation_timeout', 'request', 'target', '_predicate', '_last_yielded_row_key', '_remaining_count', '_operation_metric')

__slots__ = (
"attempt_timeout_gen",
"operation_timeout",
"request",
"target",
"_predicate",
"_last_yielded_row_key",
"_remaining_count",
"_operation_metric",
)

def __init__(
self,
query: ReadRowsQuery,
target: TargetType,
operation_timeout: float,
attempt_timeout: float,
metric: ActiveOperationMetric,
retryable_exceptions: Sequence[type[Exception]] = (),
):
self.attempt_timeout_gen = _attempt_timeout_generator(
attempt_timeout, operation_timeout
)
def __init__(self, query: ReadRowsQuery, target: TargetType, operation_timeout: float, attempt_timeout: float, metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]]=()):
self.attempt_timeout_gen = _attempt_timeout_generator(attempt_timeout, operation_timeout)
self.operation_timeout = operation_timeout
if isinstance(query, dict):
self.request = ReadRowsRequestPB(
**query, **target._request_path, app_profile_id=target.app_profile_id
)
self.request = ReadRowsRequestPB(**query, **target._request_path, app_profile_id=target.app_profile_id)
else:
self.request = query._to_pb(target)
self.target = target
Expand All @@ -108,75 +72,53 @@ def __init__(
def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]:
"""Start the read_rows operation, retrying on retryable errors.

Yields:
Row: The next row in the stream"""
return tracked_retry(
retry_fn=CrossSync._Sync_Impl.retry_target_stream,
operation=self._operation_metric,
target=self._read_rows_attempt,
predicate=self._predicate,
timeout=self.operation_timeout,
)
Yields:
Row: The next row in the stream"""
return tracked_retry(retry_fn=CrossSync._Sync_Impl.retry_target_stream, operation=self._operation_metric, target=self._read_rows_attempt, predicate=self._predicate, timeout=self.operation_timeout)

def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]:
"""Attempt a single read_rows rpc call.
This function is intended to be wrapped by retry logic,
which will call this function until it succeeds or
a non-retryable error is raised.
This function is intended to be wrapped by retry logic,
which will call this function until it succeeds or
a non-retryable error is raised.

Yields:
Row: The next row in the stream"""
Yields:
Row: The next row in the stream"""
self._operation_metric.start_attempt()
if self._last_yielded_row_key is not None:
try:
self.request.rows = self._revise_request_rowset(
row_set=self.request.rows,
last_seen_row_key=self._last_yielded_row_key,
)
self.request.rows = self._revise_request_rowset(row_set=self.request.rows, last_seen_row_key=self._last_yielded_row_key)
except _RowSetComplete:
return self.merge_rows(None)
if self._remaining_count is not None:
self.request.rows_limit = self._remaining_count
if self._remaining_count == 0:
return self.merge_rows(None)
gapic_stream = self.target.client._gapic_client.read_rows(
self.request, timeout=next(self.attempt_timeout_gen), retry=None
)
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)

def chunk_stream(
self,
stream: CrossSync._Sync_Impl.Awaitable[
CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB]
],
) -> CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk]:
def chunk_stream(self, stream: CrossSync._Sync_Impl.Awaitable[CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB]]) -> CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk]:
"""process chunks out of raw read_rows stream

Args:
stream: the raw read_rows stream from the gapic client
Yields:
ReadRowsResponsePB.CellChunk: the next chunk in the stream"""
Args:
stream: the raw read_rows stream from the gapic client
Yields:
ReadRowsResponsePB.CellChunk: the next chunk in the stream"""
current_key = None
for resp in stream:
resp = resp._pb
if resp.last_scanned_row_key:
if (
self._last_yielded_row_key is not None
and resp.last_scanned_row_key <= self._last_yielded_row_key
):
raise InvalidChunk("last scanned out of order")
if self._last_yielded_row_key is not None and resp.last_scanned_row_key <= self._last_yielded_row_key:
raise InvalidChunk('last scanned out of order')
self._last_yielded_row_key = resp.last_scanned_row_key
current_key = None
for c in resp.chunks:
if current_key is None:
current_key = c.row_key
if current_key is None:
raise InvalidChunk("first chunk is missing a row key")
elif (
self._last_yielded_row_key
and current_key <= self._last_yielded_row_key
):
raise InvalidChunk("row keys should be strictly increasing")
if not current_key:
raise InvalidChunk('first chunk is missing a row key')
elif self._last_yielded_row_key and current_key <= self._last_yielded_row_key:
raise InvalidChunk('row keys should be strictly increasing')
yield c
if c.reset_row:
current_key = None
Expand All @@ -185,18 +127,16 @@ def chunk_stream(
if self._remaining_count is not None:
self._remaining_count -= 1
if self._remaining_count < 0:
raise InvalidChunk("emit count exceeds row limit")
raise InvalidChunk('emit count exceeds row limit')
current_key = None

def merge_rows(
self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None
) -> CrossSync._Sync_Impl.Iterable[Row]:
def merge_rows(self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None) -> CrossSync._Sync_Impl.Iterable[Row]:
"""Merge chunks into rows

Args:
chunks: the chunk stream to merge
Yields:
Row: the next row in the stream"""
Args:
chunks: the chunk stream to merge
Yields:
Row: the next row in the stream"""
try:
if chunks is None:
self._operation_metric.end_with_success()
Expand All @@ -210,7 +150,7 @@ def merge_rows(
return
row_key = c.row_key
if not row_key:
raise InvalidChunk("first row chunk is missing key")
raise InvalidChunk('first row chunk is missing key')
cells = []
family: str | None = None
qualifier: bytes | None = None
Expand All @@ -220,23 +160,23 @@ def merge_rows(
raise _ResetRow(c)
k = c.row_key
f = c.family_name.value
q = c.qualifier.value if c.HasField("qualifier") else None
q = c.qualifier.value if c.HasField('qualifier') else None
if k and k != row_key:
raise InvalidChunk("unexpected new row key")
raise InvalidChunk('unexpected new row key')
if f:
family = f
if q is not None:
qualifier = q
else:
raise InvalidChunk("new family without qualifier")
raise InvalidChunk('new family without qualifier')
elif family is None:
raise InvalidChunk("missing family")
raise InvalidChunk('missing family')
elif q is not None:
if family is None:
raise InvalidChunk("new qualifier without family")
raise InvalidChunk('new qualifier without family')
qualifier = q
elif qualifier is None:
raise InvalidChunk("missing qualifier")
raise InvalidChunk('missing qualifier')
ts = c.timestamp_micros
labels = c.labels if c.labels else []
value = c.value
Expand All @@ -247,52 +187,35 @@ def merge_rows(
t = c.timestamp_micros
cl = c.labels
k = c.row_key
if (
c.HasField("family_name")
and c.family_name.value != family
):
raise InvalidChunk("family changed mid cell")
if (
c.HasField("qualifier")
and c.qualifier.value != qualifier
):
raise InvalidChunk("qualifier changed mid cell")
if c.HasField('family_name') and c.family_name.value != family:
raise InvalidChunk('family changed mid cell')
if c.HasField('qualifier') and c.qualifier.value != qualifier:
raise InvalidChunk('qualifier changed mid cell')
if t and t != ts:
raise InvalidChunk("timestamp changed mid cell")
raise InvalidChunk('timestamp changed mid cell')
if cl and cl != labels:
raise InvalidChunk("labels changed mid cell")
raise InvalidChunk('labels changed mid cell')
if k and k != row_key:
raise InvalidChunk("row key changed mid cell")
raise InvalidChunk('row key changed mid cell')
if c.reset_row:
raise _ResetRow(c)
buffer.append(c.value)
value = b"".join(buffer)
cells.append(
Cell(value, row_key, family, qualifier, ts, list(labels))
)
value = b''.join(buffer)
cells.append(Cell(value, row_key, family, qualifier, ts, list(labels)))
if c.commit_row:
block_time = time.monotonic_ns()
yield Row(row_key, cells)
if self._operation_metric.active_attempt is not None:
self._operation_metric.active_attempt.application_blocking_time_ns += (
time.monotonic_ns() - block_time
)
self._operation_metric.active_attempt.application_blocking_time_ns += time.monotonic_ns() - block_time
break
c = it.__next__()
except _ResetRow as e:
c = e.chunk
if (
c.row_key
or c.HasField("family_name")
or c.HasField("qualifier")
or c.timestamp_micros
or c.labels
or c.value
):
raise InvalidChunk("reset row with data")
if c.row_key or c.HasField('family_name') or c.HasField('qualifier') or c.timestamp_micros or c.labels or c.value:
raise InvalidChunk('reset row with data')
continue
except CrossSync._Sync_Impl.StopIteration:
raise InvalidChunk("premature end of stream")
raise InvalidChunk('premature end of stream')
except GeneratorExit as close_exception:
self._operation_metric.end_with_status(StatusCode.CANCELLED)
raise close_exception
Expand All @@ -303,19 +226,17 @@ def merge_rows(
def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB:
"""Revise the rows in the request to avoid ones we've already processed.

Args:
row_set: the row set from the request
last_seen_row_key: the last row key encountered
Returns:
RowSetPB: the new rowset after adusting for the last seen key
Raises:
_RowSetComplete: if there are no rows left to process after the revision"""
Args:
row_set: the row set from the request
last_seen_row_key: the last row key encountered
Returns:
RowSetPB: the new rowset after adusting for the last seen key
Raises:
_RowSetComplete: if there are no rows left to process after the revision"""
if row_set is None or (not row_set.row_ranges and (not row_set.row_keys)):
last_seen = last_seen_row_key
return RowSetPB(row_ranges=[RowRangePB(start_key_open=last_seen)])
adjusted_keys: list[bytes] = [
k for k in row_set.row_keys if k > last_seen_row_key
]
adjusted_keys: list[bytes] = [k for k in row_set.row_keys if k > last_seen_row_key]
adjusted_ranges: list[RowRangePB] = []
for row_range in row_set.row_ranges:
end_key = row_range.end_key_closed or row_range.end_key_open or None
Expand All @@ -327,4 +248,4 @@ def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSe
adjusted_ranges.append(new_range)
if len(adjusted_keys) == 0 and len(adjusted_ranges) == 0:
raise _RowSetComplete()
return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges)
return RowSetPB(row_keys=adjusted_keys, row_ranges=adjusted_ranges)
Loading
Loading