Skip to content

Commit

Permalink
Allow Table.read_rows to take an inclusive end key. (#3744)
Browse files Browse the repository at this point in the history
This commit adds the `end_inclusive` keyword argument, which can be explicitly passed to get `[start:end]` rather than `[start:end)`.
  • Loading branch information
lukesneeringer authored Aug 7, 2017
1 parent f034c61 commit c210494
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 6 deletions.
19 changes: 15 additions & 4 deletions bigtable/google/cloud/bigtable/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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`
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 <exceptions.ValueError>` if both
Expand All @@ -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:
Expand Down
16 changes: 14 additions & 2 deletions bigtable/tests/unit/test_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)])

Expand Down Expand Up @@ -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'
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit c210494

Please sign in to comment.