Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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(
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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:
Expand Down
24 changes: 19 additions & 5 deletions packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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

Expand Down Expand Up @@ -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):
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading