feat(bigtable): Rerouted RowSet and RowRange to use ReadRows from data client

packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py

@@ -141,6 +141,8 @@ def __repr__(self):
         return f"{self.__class__.__name__}({message!r}, {self.exceptions!r})"
 
 
+# TODO: When working on mutations batcher, rework exception handling to guarantee that
+# MutationsExceptionGroup only stores FailedMutationEntryErrors.
 class MutationsExceptionGroup(_BigtableExceptionGroup):
     """
     Represents one or more exceptions that occur during a bulk mutation operation

packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py

@@ -28,3 +28,35 @@ def batched(iterable: Iterable[T], n) -> Generator[Tuple[T, ...], None, None]:
     while batch:
         yield batch
         batch = tuple(islice(it, n))
+
+
+class _MappableAttributesMixin:
+    """
+    Mixin for classes that need some of their attribute names remapped.
+
+    This is for taking some of the classes from the data client row filters
+    and row range classes that are 1:1 with their legacy client counterparts but with
+    some of their attributes renamed. To use in a class, override the base class with this mixin
+    class and define a map _attribute_map from legacy client attributes to data client
+    attributes.
+
+    Attributes are remapped and redefined in __init__ as well as getattr/setattr.
+    """
+
+    def __init__(self, *args, **kwargs):
+        new_kwargs = {self._attribute_map.get(k, k): v for (k, v) in kwargs.items()}
+        super(_MappableAttributesMixin, self).__init__(*args, **new_kwargs)
+
+    def __getattr__(self, name):
+        if name == "_attribute_map":
+            raise AttributeError
+        if name not in self._attribute_map:
+            raise AttributeError
+        return getattr(self, self._attribute_map[name])
+
+    def __setattr__(self, name, value):
+        if name == "_attribute_map":
+            super(_MappableAttributesMixin, self).__setattr__(name, value)
+            return
+        attribute = self._attribute_map.get(name, name)
+        super(_MappableAttributesMixin, self).__setattr__(attribute, value)

packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py

@@ -50,6 +50,7 @@
 from google.cloud.bigtable.data.row_filters import (
     TimestampRangeFilter as BaseTimestampRangeFilter,
 )
+from google.cloud.bigtable.helpers import _MappableAttributesMixin
 
 _PACK_I64 = struct.Struct(">q").pack
 
@@ -70,35 +71,6 @@ def _deprecated_to_pb(self):
 # Provide to_pb alias with DeprecationWarning for backwards compatibility with legacy client code
 RowFilter.to_pb = _deprecated_to_pb  # type: ignore[attr-defined]
 TimestampRange.to_pb = _deprecated_to_pb  # type: ignore[attr-defined]
-
-
-class _MappableAttributesMixin:
-    """
-    Mixin for classes that need some of their attribute names remapped.
-
-    This is for taking some of the classes from the data client row filters
-    that are 1:1 with their legacy client counterparts but with some of their
-    attributes renamed. To use in a class, override the base class with this mixin
-    class and define a map _attribute_map from legacy client attributes to data client
-    attributes.
-
-    Attributes are remapped and redefined in __init__ as well as getattr/setattr.
-    """
-
-    def __init__(self, *args, **kwargs):
-        new_kwargs = {self._attribute_map.get(k, k): v for (k, v) in kwargs.items()}
-        super(_MappableAttributesMixin, self).__init__(*args, **new_kwargs)
-
-    def __getattr__(self, name):
-        if name not in self._attribute_map:
-            raise AttributeError
-        return getattr(self, self._attribute_map[name])
-
-    def __setattr__(self, name, value):
-        attribute = self._attribute_map.get(name, name)
-        super(_MappableAttributesMixin, self).__setattr__(attribute, value)
-
-
 # The classes defined below are to provide constructors and members
 # that have an interface that does not match the one used by the data
 # client, for backwards compatibility purposes.

packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py

@@ -14,7 +14,15 @@
 
 """User-friendly container for Google Cloud Bigtable RowSet"""
 
-from google.cloud._helpers import _to_bytes  # type: ignore
+from google.cloud._helpers import _to_bytes
+
+from google.cloud.bigtable.data.read_rows_query import (
+    ReadRowsQuery,
+)
+from google.cloud.bigtable.data.read_rows_query import (
+    RowRange as BaseRowRange,
+)
+from google.cloud.bigtable.helpers import _MappableAttributesMixin
 
 
 class RowSet(object):
@@ -25,30 +33,25 @@ class RowSet(object):
     """
 
     def __init__(self):
-        self.row_keys = []
-        self.row_ranges = []
+        self._read_rows_query = ReadRowsQuery()
 
     def __eq__(self, other):
         if not isinstance(other, self.__class__):
             return NotImplemented
 
-        if len(other.row_keys) != len(self.row_keys):
-            return False
-
-        if len(other.row_ranges) != len(self.row_ranges):
-            return False
-
-        if not set(other.row_keys) == set(self.row_keys):
-            return False
-
-        if not set(other.row_ranges) == set(self.row_ranges):
-            return False
-
-        return True
+        return self._read_rows_query == other._read_rows_query
 
     def __ne__(self, other):
         return not self == other
 
+    @property
+    def row_keys(self):
+        return self._read_rows_query.row_keys
+
+    @property
+    def row_ranges(self):
+        return self._read_rows_query.row_ranges
+
     def add_row_key(self, row_key):
         """Add row key to row_keys list.
 
@@ -62,7 +65,7 @@ def add_row_key(self, row_key):
         :type row_key: bytes
         :param row_key: The key of a row to read
         """
-        self.row_keys.append(row_key)
+        self._read_rows_query.add_key(row_key)
 
     def add_row_range(self, row_range):
         """Add row_range to row_ranges list.
@@ -77,7 +80,7 @@ def add_row_range(self, row_range):
         :type row_range: class:`RowRange`
         :param row_range: The row range object having start and end key
         """
-        self.row_ranges.append(row_range)
+        self._read_rows_query.add_range(row_range)
 
     def add_row_range_from_keys(
         self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False
@@ -109,7 +112,7 @@ def add_row_range_from_keys(
                   considered inclusive. The default is False (exclusive).
         """
         row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive)
-        self.row_ranges.append(row_range)
+        self._read_rows_query.add_range(row_range)
 
     def add_row_range_with_prefix(self, row_key_prefix):
         """Add row range to row_ranges list that start with the row_key_prefix from the row keys
@@ -135,22 +138,21 @@ def _update_message_request(self, message):
         :type message: class:`data_messages_v2_pb2.ReadRowsRequest`
         :param message: The ``ReadRowsRequest`` protobuf
         """
-        for each in self.row_keys:
+        for each in self._read_rows_query.row_keys:
             message.rows.row_keys._pb.append(_to_bytes(each))
 
-        for each in self.row_ranges:
-            r_kwrags = each.get_range_kwargs()
-            message.rows.row_ranges.append(r_kwrags)
+        for each in self._read_rows_query.row_ranges:
+            message.rows.row_ranges.append(each._to_pb())
 
 
-class RowRange(object):
+class RowRange(_MappableAttributesMixin, BaseRowRange):
     """Convenience wrapper of google.bigtable.v2.RowRange
 
-    :type start_key: bytes
+    :type start_key: str | bytes
     :param start_key: (Optional) Start key of the row range. If left empty,
                       will be interpreted as the empty string.
 
-    :type end_key: bytes
+    :type end_key: str | bytes
     :param end_key: (Optional) End key of the row range. If left empty,
                     will be interpreted as the empty string and range will
                     be unbounded on the high end.
@@ -164,49 +166,31 @@ class RowRange(object):
                   considered inclusive. The default is False (exclusive).
     """
 
-    def __init__(
-        self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False
-    ):
-        self.start_key = start_key
-        self.start_inclusive = start_inclusive
-        self.end_key = end_key
-        self.end_inclusive = end_inclusive
+    _attribute_map = {
+        "start_inclusive": "start_is_inclusive",
+        "end_inclusive": "end_is_inclusive",
+    }
 
     def _key(self):
-        """A tuple key that uniquely describes this field.
-
-        Used to compute this instance's hashcode and evaluate equality.
-
-        Returns:
-            Tuple[str]: The contents of this :class:`.RowRange`.
-        """
-        return (self.start_key, self.start_inclusive, self.end_key, self.end_inclusive)
+        return (
+            self.start_key,
+            self.end_key,
+            self.start_is_inclusive,
+            self.end_is_inclusive,
+        )
 
     def __hash__(self):
         return hash(self._key())
 
-    def __eq__(self, other):
-        if not isinstance(other, self.__class__):
-            return NotImplemented
-        return self._key() == other._key()
-
-    def __ne__(self, other):
-        return not self == other
-
     def get_range_kwargs(self):
         """Convert row range object to dict which can be passed to
         google.bigtable.v2.RowRange add method.
         """
         range_kwargs = {}
         if self.start_key is not None:
-            start_key_key = "start_key_open"
-            if self.start_inclusive:
-                start_key_key = "start_key_closed"
-            range_kwargs[start_key_key] = _to_bytes(self.start_key)
-
+            key = "start_key_closed" if self.start_is_inclusive else "start_key_open"
+            range_kwargs[key] = self.start_key
         if self.end_key is not None:
-            end_key_key = "end_key_open"
-            if self.end_inclusive:
-                end_key_key = "end_key_closed"
-            range_kwargs[end_key_key] = _to_bytes(self.end_key)
+            key = "end_key_closed" if self.end_is_inclusive else "end_key_open"
+            range_kwargs[key] = self.end_key
         return range_kwargs

packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py

@@ -48,38 +48,6 @@ def test_row_set__eq__type_differ():
     assert not (row_set1 == row_set2)
 
 
-def test_row_set__eq__len_row_keys_differ():
-    from google.cloud.bigtable.row_set import RowSet
-
-    row_key1 = b"row_key1"
-    row_key2 = b"row_key1"
-
-    row_set1 = RowSet()
-    row_set2 = RowSet()
-
-    row_set1.add_row_key(row_key1)
-    row_set1.add_row_key(row_key2)
-    row_set2.add_row_key(row_key2)
-
-    assert not (row_set1 == row_set2)
-
-
-def test_row_set__eq__len_row_ranges_differ():
-    from google.cloud.bigtable.row_set import RowRange, RowSet
-
-    row_range1 = RowRange(b"row_key4", b"row_key9")
-    row_range2 = RowRange(b"row_key4", b"row_key9")
-
-    row_set1 = RowSet()
-    row_set2 = RowSet()
-
-    row_set1.add_row_range(row_range1)
-    row_set1.add_row_range(row_range2)
-    row_set2.add_row_range(row_range2)
-
-    assert not (row_set1 == row_set2)
-
-
 def test_row_set__eq__row_keys_differ():
     from google.cloud.bigtable.row_set import RowSet
 
@@ -223,8 +191,8 @@ def test_row_range_constructor():
     start_key = "row_key1"
     end_key = "row_key9"
     row_range = RowRange(start_key, end_key)
-    assert start_key == row_range.start_key
-    assert end_key == row_range.end_key
+    assert start_key == row_range.start_key.decode()
+    assert end_key == row_range.end_key.decode()
     assert row_range.start_inclusive
     assert not row_range.end_inclusive
 
@@ -311,6 +279,68 @@ def test_row_range_get_range_kwargs_open_closed():
     assert expected_result == actual_result
 
 
+def test_row_range_get_range_kwargs_unbounded_start():
+    from google.cloud.bigtable.row_set import RowRange
+
+    end_key = b"row_key9"
+    expected_result = {"end_key_open": end_key}
+    row_range = RowRange(None, end_key)
+    actual_result = row_range.get_range_kwargs()
+    assert expected_result == actual_result
+
+
+def test_row_range_get_range_kwargs_unbounded_end():
+    from google.cloud.bigtable.row_set import RowRange
+
+    start_key = b"row_key1"
+    expected_result = {"start_key_closed": start_key}
+    row_range = RowRange(start_key, None)
+    actual_result = row_range.get_range_kwargs()
+    assert expected_result == actual_result
+
+
+def test_mappable_attributes_mixin_guard_recursion():
+    import pytest
+
+    from google.cloud.bigtable.helpers import _MappableAttributesMixin
+
+    class Uninitialized(_MappableAttributesMixin):
+        pass
+
+    obj = Uninitialized()
+    with pytest.raises(AttributeError):
+        _ = obj._attribute_map
+    with pytest.raises(AttributeError):
+        _ = obj.non_existent_attribute
+
+    # Verify setattr on _attribute_map doesn't cause recursion
+    obj._attribute_map = {"old_attr": "new_attr"}
+    assert obj._attribute_map == {"old_attr": "new_attr"}
+
+
+def test_mappable_attributes_mixin_remap():
+    import pytest
+
+    from google.cloud.bigtable.helpers import _MappableAttributesMixin
+
+    class DummyBase:
+        def __init__(self, target_attr=None):
+            self.target_attr = target_attr
+
+    class MappedClass(_MappableAttributesMixin, DummyBase):
+        _attribute_map = {"source_attr": "target_attr"}
+
+    instance = MappedClass(source_attr="hello")
+    assert instance.target_attr == "hello"
+    assert instance.source_attr == "hello"
+
+    instance.source_attr = "world"
+    assert instance.target_attr == "world"
+
+    with pytest.raises(AttributeError):
+        _ = instance.unknown_attr
+
+
 def _ReadRowsRequestPB(*args, **kw):
     from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2