fix: improve rowset revision (#979)

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

@@ -342,7 +342,7 @@ def _revise_request_rowset(
             _RowSetComplete: if there are no rows left to process after the revision
         """
         # if user is doing a whole table scan, start a new one with the last seen key
-        if row_set is None or (not row_set.row_ranges and row_set.row_keys is not None):
+        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)])
         # remove seen keys from user-specific key list

tests/unit/data/_async/test__read_rows.py

@@ -98,19 +98,31 @@ def test_ctor(self):
             (["d", "c", "b", "a"], "b", ["d", "c"]),
         ],
     )
-    def test_revise_request_rowset_keys(self, in_keys, last_key, expected):
+    @pytest.mark.parametrize("with_range", [True, False])
+    def test_revise_request_rowset_keys_with_range(
+        self, in_keys, last_key, expected, with_range
+    ):
         from google.cloud.bigtable_v2.types import RowSet as RowSetPB
         from google.cloud.bigtable_v2.types import RowRange as RowRangePB
+        from google.cloud.bigtable.data.exceptions import _RowSetComplete
 
         in_keys = [key.encode("utf-8") for key in in_keys]
         expected = [key.encode("utf-8") for key in expected]
         last_key = last_key.encode("utf-8")
 
-        sample_range = RowRangePB(start_key_open=last_key)
-        row_set = RowSetPB(row_keys=in_keys, row_ranges=[sample_range])
-        revised = self._get_target_class()._revise_request_rowset(row_set, last_key)
-        assert revised.row_keys == expected
-        assert revised.row_ranges == [sample_range]
+        if with_range:
+            sample_range = [RowRangePB(start_key_open=last_key)]
+        else:
+            sample_range = []
+        row_set = RowSetPB(row_keys=in_keys, row_ranges=sample_range)
+        if not with_range and expected == []:
+            # expect exception if we are revising to an empty rowset
+            with pytest.raises(_RowSetComplete):
+                self._get_target_class()._revise_request_rowset(row_set, last_key)
+        else:
+            revised = self._get_target_class()._revise_request_rowset(row_set, last_key)
+            assert revised.row_keys == expected
+            assert revised.row_ranges == sample_range
 
     @pytest.mark.parametrize(
         "in_ranges,last_key,expected",
@@ -157,9 +169,13 @@ def test_revise_request_rowset_keys(self, in_keys, last_key, expected):
             ),
         ],
     )
-    def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected):
+    @pytest.mark.parametrize("with_key", [True, False])
+    def test_revise_request_rowset_ranges(
+        self, in_ranges, last_key, expected, with_key
+    ):
         from google.cloud.bigtable_v2.types import RowSet as RowSetPB
         from google.cloud.bigtable_v2.types import RowRange as RowRangePB
+        from google.cloud.bigtable.data.exceptions import _RowSetComplete
 
         # convert to protobuf
         next_key = (last_key + "a").encode("utf-8")
@@ -172,10 +188,20 @@ def test_revise_request_rowset_ranges(self, in_ranges, last_key, expected):
             RowRangePB(**{k: v.encode("utf-8") for k, v in r.items()}) for r in expected
         ]
 
-        row_set = RowSetPB(row_ranges=in_ranges, row_keys=[next_key])
-        revised = self._get_target_class()._revise_request_rowset(row_set, last_key)
-        assert revised.row_keys == [next_key]
-        assert revised.row_ranges == expected
+        if with_key:
+            row_keys = [next_key]
+        else:
+            row_keys = []
+
+        row_set = RowSetPB(row_ranges=in_ranges, row_keys=row_keys)
+        if not with_key and expected == []:
+            # expect exception if we are revising to an empty rowset
+            with pytest.raises(_RowSetComplete):
+                self._get_target_class()._revise_request_rowset(row_set, last_key)
+        else:
+            revised = self._get_target_class()._revise_request_rowset(row_set, last_key)
+            assert revised.row_keys == row_keys
+            assert revised.row_ranges == expected
 
     @pytest.mark.parametrize("last_key", ["a", "b", "c"])
     def test_revise_request_full_table(self, last_key):