diff --git a/src/google/cloud/happybase/table.py b/src/google/cloud/happybase/table.py index b9c0715..8824cf2 100644 --- a/src/google/cloud/happybase/table.py +++ b/src/google/cloud/happybase/table.py @@ -35,6 +35,7 @@ from google.cloud.bigtable.row_filters import TimestampRange from google.cloud.bigtable.row_filters import TimestampRangeFilter from google.cloud.bigtable.table import Table as _LowLevelTable +from google.cloud.bigtable.row_set import RowSet from google.cloud.happybase.batch import _get_column_pairs from google.cloud.happybase.batch import _WAL_SENTINEL @@ -235,7 +236,7 @@ def rows(self, rows, columns=None, timestamp=None, # versions == 1 since we only want the latest. filter_ = _filter_chain_helper(versions=1, timestamp=timestamp, filters=filters) - rows_generator = self._low_level_table.yield_rows( + rows_generator = self._low_level_table.read_rows( filter_=filter_) # NOTE: We could use max_loops = 1000 or some similar value to ensure # that the stream isn't open too long. @@ -381,9 +382,10 @@ def scan(self, row_start=None, row_stop=None, row_prefix=None, row_start, row_stop, filter_chain = _scan_filter_helper( row_start, row_stop, row_prefix, columns, timestamp, limit, kwargs) - rows_generator = self._low_level_table.yield_rows( - start_key=row_start, end_key=row_stop, - limit=limit, filter_=filter_chain) + row_set = _get_row_set_object(row_start, row_stop) + + rows_generator = self._low_level_table.read_rows( + row_set=row_set, limit=limit, filter_=filter_chain) for rowdata in rows_generator: curr_row_data = rowdata @@ -969,3 +971,12 @@ def _row_keys_filter_helper(row_keys): return filters[0] else: return RowFilterUnion(filters=filters) + + +def _get_row_set_object(row_start, row_stop): + """Return a RowSet object for the given row_start and row_stop + """ + row_set = RowSet() + row_set.add_row_range_from_keys(start_key=row_start, + end_key=row_stop) + return row_set diff --git a/unit_tests/test_table.py b/unit_tests/test_table.py index 8edf317..3bcf857 100644 --- a/unit_tests/test_table.py +++ b/unit_tests/test_table.py @@ -619,8 +619,14 @@ def mock_filter_chain_helper(**kwargs): mock_filters.append(kwargs) return fake_filter + fake_row_set = object() + + def mock_get_row_set_object(*args): # pylint: disable=unused-argument + return fake_row_set + with _Monkey(MUT, _filter_chain_helper=mock_filter_chain_helper, - _columns_filter_helper=mock_columns_filter_helper): + _columns_filter_helper=mock_columns_filter_helper, + _get_row_set_object=mock_get_row_set_object): result = table.scan(row_start=row_start, row_stop=row_stop, row_prefix=row_prefix, columns=columns, filter=filter_, timestamp=timestamp, @@ -636,12 +642,13 @@ def mock_filter_chain_helper(**kwargs): if row_prefix: row_start = row_prefix row_stop = MUT._string_successor(row_prefix) + read_rows_kwargs = { - 'end_key': row_stop, + 'row_set': fake_row_set, 'filter_': fake_filter, 'limit': limit, - 'start_key': row_start, } + self.assertEqual(table._low_level_table.read_rows_calls, [ (read_rows_args, read_rows_kwargs), ]) @@ -1427,6 +1434,22 @@ def test_many_rows(self): self.assertEqual(result, expected_result) +class Test___get_row_set_object(unittest.TestCase): + + def _callFUT(self, *args, **kwargs): + from google.cloud.happybase.table import _get_row_set_object + return _get_row_set_object(*args, **kwargs) + + def test_row_set_object(self): + from google.cloud.bigtable.row_set import RowSet + + start_key = b"row_key2" + end_key = b"row_key9" + + row_set = self._callFUT(start_key, end_key) + self.assertIsInstance(row_set, RowSet) + + class _Connection(object): def __init__(self, instance): @@ -1466,7 +1489,7 @@ def read_row(self, *args, **kwargs): self.read_row_calls.append((args, kwargs)) return self.read_row_result - def yield_rows(self, *args, **kwargs): + def read_rows(self, *args, **kwargs): self.read_rows_calls.append((args, kwargs)) rows_dict = self.read_rows_result.rows for row_key in sorted(rows_dict):