Allow Table.read_rows to take an inclusive end key. (googleapis#3744)

This commit adds the `end_inclusive` keyword argument, which can be explicitly passed to get `[start:end]` rather than `[start:end)`.

Author: lukesneeringer

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 <exceptions.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:

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