diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py index 40ef3a2ca2fb9..64fbcc93771ed 100644 --- a/bigtable/google/cloud/bigtable/table.py +++ b/bigtable/google/cloud/bigtable/table.py @@ -257,7 +257,7 @@ def read_row(self, row_key, filter_=None): return rows_data.rows[row_key] def read_rows(self, start_key=None, end_key=None, limit=None, - filter_=None): + filter_=None, end_inclusive=False): """Read rows from this table. :type start_key: bytes @@ -280,13 +280,17 @@ def read_rows(self, start_key=None, end_key=None, limit=None, specified row(s). If unset, reads every column in each row. + :type end_inclusive: bool + :param end_inclusive: (Optional) Whether the ``end_key`` should be + considered inclusive. The default is False (exclusive). + :rtype: :class:`.PartialRowsData` :returns: A :class:`.PartialRowsData` convenience wrapper for consuming the streamed results. """ request_pb = _create_row_request( self.name, start_key=start_key, end_key=end_key, filter_=filter_, - limit=limit) + limit=limit, end_inclusive=end_inclusive) client = self._instance._client response_iterator = client._data_stub.ReadRows(request_pb) # We expect an iterator of `data_messages_v2_pb2.ReadRowsResponse` @@ -360,7 +364,7 @@ def sample_row_keys(self): def _create_row_request(table_name, row_key=None, start_key=None, end_key=None, - filter_=None, limit=None): + filter_=None, limit=None, end_inclusive=False): """Creates a request to read rows in a table. :type table_name: str @@ -388,6 +392,10 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None, rows' worth of results. The default (zero) is to return all results. + :type end_inclusive: bool + :param end_inclusive: (Optional) Whether the ``end_key`` should be + considered inclusive. The default is False (exclusive). + :rtype: :class:`data_messages_v2_pb2.ReadRowsRequest` :returns: The ``ReadRowsRequest`` protobuf corresponding to the inputs. :raises: :class:`ValueError ` if both @@ -403,7 +411,10 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None, if start_key is not None: range_kwargs['start_key_closed'] = _to_bytes(start_key) if end_key is not None: - range_kwargs['end_key_open'] = _to_bytes(end_key) + end_key_key = 'end_key_open' + if end_inclusive: + end_key_key = 'end_key_closed' + range_kwargs[end_key_key] = _to_bytes(end_key) if filter_ is not None: request_kwargs['filter'] = filter_.to_pb() if limit is not None: diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py index dc4d2b5bbad08..3890d097f572f 100644 --- a/bigtable/tests/unit/test_table.py +++ b/bigtable/tests/unit/test_table.py @@ -537,6 +537,7 @@ def mock_create_row_request(table_name, **kwargs): 'end_key': end_key, 'filter_': filter_obj, 'limit': limit, + 'end_inclusive': False, } self.assertEqual(mock_created, [(table.name, created_kwargs)]) @@ -572,12 +573,12 @@ def test_sample_row_keys(self): class Test__create_row_request(unittest.TestCase): def _call_fut(self, table_name, row_key=None, start_key=None, end_key=None, - filter_=None, limit=None): + filter_=None, limit=None, end_inclusive=False): from google.cloud.bigtable.table import _create_row_request return _create_row_request( table_name, row_key=row_key, start_key=start_key, end_key=end_key, - filter_=filter_, limit=limit) + filter_=filter_, limit=limit, end_inclusive=end_inclusive) def test_table_name_only(self): table_name = 'table_name' @@ -627,6 +628,17 @@ def test_row_range_both_keys(self): start_key_closed=start_key, end_key_open=end_key) self.assertEqual(result, expected_result) + def test_row_range_both_keys_inclusive(self): + table_name = 'table_name' + start_key = b'start_key' + end_key = b'end_key' + result = self._call_fut(table_name, start_key=start_key, + end_key=end_key, end_inclusive=True) + expected_result = _ReadRowsRequestPB(table_name=table_name) + expected_result.rows.row_ranges.add( + start_key_closed=start_key, end_key_closed=end_key) + self.assertEqual(result, expected_result) + def test_with_filter(self): from google.cloud.bigtable.row_filters import RowSampleFilter