diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py index d6fcfc73520b..fad50c976946 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py @@ -53,6 +53,27 @@ __CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen._read_rows" +def _get_pb_range_bounds( + row_range: RowRangePB, +) -> tuple[bytes | None, bytes | None]: + """Returns (start_key, end_key) from a RowRange protobuf, or None if unbounded.""" + start_key = ( + row_range.start_key_closed + if "start_key_closed" in row_range + else row_range.start_key_open + if "start_key_open" in row_range + else None + ) + end_key = ( + row_range.end_key_closed + if "end_key_closed" in row_range + else row_range.end_key_open + if "end_key_open" in row_range + else None + ) + return start_key, end_key + + @CrossSync.convert_class("_ReadRowsOperation") class _ReadRowsOperationAsync: """ @@ -377,11 +398,10 @@ def _revise_request_rowset( # adjust ranges to ignore keys before last seen 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 + start_key, end_key = _get_pb_range_bounds(row_range) if end_key is None or end_key > last_seen_row_key: # end range is after last seen key new_range = RowRangePB(row_range) - start_key = row_range.start_key_closed or row_range.start_key_open if start_key is None or start_key <= last_seen_row_key: # replace start key with last seen new_range.start_key_open = last_seen_row_key diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index f8c2d28b2903..9788f67e4f5d 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -48,6 +48,27 @@ ) +def _get_pb_range_bounds( + row_range: RowRangePB, +) -> tuple[bytes | None, bytes | None]: + """Returns (start_key, end_key) from a RowRange protobuf, or None if unbounded.""" + start_key = ( + row_range.start_key_closed + if "start_key_closed" in row_range + else row_range.start_key_open + if "start_key_open" in row_range + else None + ) + end_key = ( + row_range.end_key_closed + if "end_key_closed" in row_range + else row_range.end_key_open + if "end_key_open" in row_range + else None + ) + return start_key, end_key + + class _ReadRowsOperation: """ ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream @@ -318,10 +339,9 @@ def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSe ] 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 + start_key, end_key = _get_pb_range_bounds(row_range) if end_key is None or end_key > last_seen_row_key: new_range = RowRangePB(row_range) - start_key = row_range.start_key_closed or row_range.start_key_open if start_key is None or start_key <= last_seen_row_key: new_range.start_key_open = last_seen_row_key adjusted_ranges.append(new_range) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py index 938303ea5a4c..2fc831950e25 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/read_rows_query.py @@ -52,16 +52,7 @@ def __init__( start_is_inclusive: bool | None = None, end_is_inclusive: bool | None = None, ): - # convert empty key inputs to None for consistency - start_key = None if not start_key else start_key - end_key = None if not end_key else end_key - # check for invalid combinations of arguments - if start_is_inclusive is None: - start_is_inclusive = True - - if end_is_inclusive is None: - end_is_inclusive = False - # ensure that start_key and end_key are bytes + # ensure that start_key and end_key are bytes or None if isinstance(start_key, str): start_key = start_key.encode() elif start_key is not None and not isinstance(start_key, bytes): @@ -70,21 +61,39 @@ def __init__( end_key = end_key.encode() elif end_key is not None and not isinstance(end_key, bytes): raise ValueError("end_key must be a string or bytes") + + # Record whether end_is_inclusive was explicitly specified before applying defaults + end_inclusive_omitted = end_is_inclusive is None + + # Apply default inclusivity + if start_is_inclusive is None: + start_is_inclusive = True + if end_is_inclusive is None: + end_is_inclusive = False + + # Normalize default empty inputs (e.g. RowRange("", "")) to unbounded None: + # - [b"", ...) covers all valid (non-empty) Bigtable rows, identical to [-inf, ...). + # Only preserve b"" if the caller explicitly requested exclusive start (start_is_inclusive=False). + if start_key == b"" and start_is_inclusive: + start_key = None + # - An empty end_key was historically passed as shorthand for unbounded (+inf) only + # when end_is_inclusive was omitted (e.g. RowRange(start, "")). + if end_key == b"" and end_inclusive_omitted: + end_key = None + # ensure that start_key is less than or equal to end_key if start_key is not None and end_key is not None and start_key > end_key: raise ValueError("start_key must be less than or equal to end_key") init_dict = {} if start_key is not None: - if start_is_inclusive: - init_dict["start_key_closed"] = start_key - else: - init_dict["start_key_open"] = start_key + init_dict[ + "start_key_closed" if start_is_inclusive else "start_key_open" + ] = start_key if end_key is not None: - if end_is_inclusive: - init_dict["end_key_closed"] = end_key - else: - init_dict["end_key_open"] = end_key + init_dict["end_key_closed" if end_is_inclusive else "end_key_open"] = ( + end_key + ) self._pb = RowRangePB(**init_dict) def _key( @@ -105,7 +114,11 @@ def start_key(self) -> bytes | None: """ Returns the start key of the range. If None, the range is unbounded on the left. """ - return self._pb.start_key_closed or self._pb.start_key_open or None + if "start_key_closed" in self._pb: + return self._pb.start_key_closed + if "start_key_open" in self._pb: + return self._pb.start_key_open + return None @property def end_key(self) -> bytes | None: @@ -115,7 +128,11 @@ def end_key(self) -> bytes | None: Returns: bytes | None: The end key of the range, or None if the range is unbounded on the right. """ - return self._pb.end_key_closed or self._pb.end_key_open or None + if "end_key_closed" in self._pb: + return self._pb.end_key_closed + if "end_key_open" in self._pb: + return self._pb.end_key_open + return None @property def start_is_inclusive(self) -> bool: @@ -127,7 +144,7 @@ def start_is_inclusive(self) -> bool: Returns: bool: Whether the range is inclusive of the start key. """ - return not bool(self._pb.start_key_open) + return "start_key_open" not in self._pb @property def end_is_inclusive(self) -> bool: @@ -139,7 +156,7 @@ def end_is_inclusive(self) -> bool: Returns: bool: Whether the range is inclusive of the end key. """ - return not bool(self._pb.end_key_open) + return "end_key_open" not in self._pb def _to_pb(self) -> RowRangePB: """ @@ -190,10 +207,10 @@ def __bool__(self) -> bool: bool: True if the RowRange is not empty, False otherwise """ return bool( - self._pb.start_key_closed - or self._pb.start_key_open - or self._pb.end_key_closed - or self._pb.end_key_open + "start_key_closed" in self._pb + or "start_key_open" in self._pb + or "end_key_closed" in self._pb + or "end_key_open" in self._pb ) def __eq__(self, other: Any) -> bool: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py index 6e041e3ed7b8..077c7b9b2c43 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py @@ -14,6 +14,8 @@ """User-friendly container for Google Cloud Bigtable RowSet""" +from google.cloud._helpers import _to_bytes # type: ignore + from google.cloud.bigtable.data.read_rows_query import ( ReadRowsQuery, ) @@ -23,6 +25,17 @@ from google.cloud.bigtable.helpers import _MappableAttributesMixin +def _calculate_prefix_end_key(prefix: bytes) -> bytes | None: + """Calculates the exclusive upper bound for a prefix scan by incrementing the last byte. + + Trailing 0xFF bytes roll over. Returns None if prefix consists entirely of 0xFF (no upper bound). + """ + stripped = prefix.rstrip(b"\xff") + if not stripped: + return None + return stripped[:-1] + bytes([stripped[-1] + 1]) + + class RowSet(object): """Convenience wrapper of google.bigtable.v2.RowSet @@ -121,14 +134,15 @@ def add_row_range_with_prefix(self, row_key_prefix): :start-after: [START bigtable_api_add_row_range_with_prefix] :end-before: [END bigtable_api_add_row_range_with_prefix] - :type row_key_prefix: str + :type row_key_prefix: str | bytes :param row_key_prefix: To retrieve all rows that start with this row key prefix. Prefix cannot be zero length.""" - end_key = row_key_prefix[:-1] + chr(ord(row_key_prefix[-1]) + 1) - self.add_row_range_from_keys( - row_key_prefix.encode("utf-8"), end_key.encode("utf-8") - ) + prefix_bytes = _to_bytes(row_key_prefix, encoding="utf-8") + if not prefix_bytes: + raise ValueError("row_key_prefix cannot be empty.") + end_key = _calculate_prefix_end_key(prefix_bytes) + self.add_row_range_from_keys(prefix_bytes, end_key) class RowRange(_MappableAttributesMixin, BaseRowRange): diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py index 24b92d2de5ad..96387707244d 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py @@ -160,6 +160,13 @@ def test_revise_request_rowset_keys_with_range( "a", [{"start_key_open": "a", "end_key_open": "z"}], ), + ([{"end_key_closed": ""}], "", []), + ([{"end_key_open": ""}], "", []), + ( + [{"start_key_open": "", "end_key_closed": "d"}], + "", + [{"start_key_open": "", "end_key_closed": "d"}], + ), ], ) @pytest.mark.parametrize("with_key", [True, False]) diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py index b6cac9cafbc0..511d20794d80 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py @@ -155,6 +155,13 @@ def test_revise_request_rowset_keys_with_range( "a", [{"start_key_open": "a", "end_key_open": "z"}], ), + ([{"end_key_closed": ""}], "", []), + ([{"end_key_open": ""}], "", []), + ( + [{"start_key_open": "", "end_key_closed": "d"}], + "", + [{"start_key_open": "", "end_key_closed": "d"}], + ), ], ) @pytest.mark.parametrize("with_key", [True, False]) diff --git a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py index cfef2da4b793..28b31e5f0dda 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test_read_rows_query.py @@ -57,11 +57,40 @@ def test_ctor_empty_strings(self): """ empty strings should be treated as None """ - row_range = self._make_one("", "") - assert row_range.start_key is None - assert row_range.end_key is None - assert row_range.start_is_inclusive is True - assert row_range.end_is_inclusive is True + from google.cloud.bigtable_v2.types import RowRange as RowRangePB + + for empty in ("", b""): + row_range = self._make_one(empty, empty) + assert row_range.start_key is None + assert row_range.end_key is None + assert row_range.start_is_inclusive is True + assert row_range.end_is_inclusive is True + + rr_open_start = self._make_one(start_key=b"", start_is_inclusive=False) + assert rr_open_start.start_key == b"" + assert rr_open_start.start_is_inclusive is False + assert bool(rr_open_start) is True + assert "start_key_open" in rr_open_start._to_pb() + + rr_closed_end = self._make_one(end_key=b"", end_is_inclusive=True) + assert rr_closed_end.end_key == b"" + assert rr_closed_end.end_is_inclusive is True + assert bool(rr_closed_end) is True + assert "end_key_closed" in rr_closed_end._to_pb() + + rr_open_end = self._make_one(end_key=b"", end_is_inclusive=False) + assert rr_open_end.end_key == b"" + assert rr_open_end.end_is_inclusive is False + assert bool(rr_open_end) is True + assert "end_key_open" in rr_open_end._to_pb() + + rr_from_pb = self._get_target_class()._from_pb( + RowRangePB(start_key_open=b"", end_key_closed=b"z") + ) + assert rr_from_pb.start_key == b"" + assert rr_from_pb.start_is_inclusive is False + assert rr_from_pb.end_key == b"z" + assert rr_from_pb.end_is_inclusive is True def test_ctor_inclusive_flags(self): row_range = self._make_one("test_row5", "test_row6", False, True) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py index e76d1f15135c..5ed3ce223364 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py @@ -163,6 +163,33 @@ def test_row_set_add_row_range_with_prefix(): assert row_set.row_ranges[0].end_key == b"rox" +def test_row_set_add_row_range_with_prefix_utf8_boundaries(): + import pytest + + from google.cloud.bigtable.row_set import RowSet + + cases = [ + ("user\x7f", b"user\x7f", b"user\x80"), + ("p\u07ff", b"p\xdf\xbf", b"p\xdf\xc0"), + ("p\uffff", b"p\xef\xbf\xbf", b"p\xef\xbf\xc0"), + ("prefix\ud7ff", b"prefix\xed\x9f\xbf", b"prefix\xed\x9f\xc0"), + ("prefix\U0010ffff", b"prefix\xf4\x8f\xbf\xbf", b"prefix\xf4\x8f\xbf\xc0"), + (b"abc\xff", b"abc\xff", b"abd"), + (b"\xff\xff", b"\xff\xff", None), + ] + for prefix, expected_start, expected_end in cases: + row_set = RowSet() + row_set.add_row_range_with_prefix(prefix) + assert row_set.row_ranges[0].start_key == expected_start + assert row_set.row_ranges[0].end_key == expected_end + + row_set = RowSet() + with pytest.raises(ValueError): + row_set.add_row_range_with_prefix("") + with pytest.raises(ValueError): + row_set.add_row_range_with_prefix(b"") + + def test_row_range_constructor(): from google.cloud.bigtable.row_set import RowRange