Merge pull request #1321 from dhermes/bigtable-row-filter-8

Implementing Bigtable TimestampRangeFilter.

Author: dhermes

gcloud/bigtable/row.py

@@ -15,6 +15,7 @@
 """User friendly container for Google Cloud Bigtable Row."""
 
 
+from gcloud._helpers import _microseconds_from_datetime
 from gcloud._helpers import _to_bytes
 from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
 
@@ -252,6 +253,74 @@ def to_pb(self):
         return data_pb2.RowFilter(column_qualifier_regex_filter=self.regex)
 
 
+class TimestampRange(object):
+    """Range of time with inclusive lower and exclusive upper bounds.
+
+    :type start: :class:`datetime.datetime`
+    :param start: (Optional) The (inclusive) lower bound of the timestamp
+                  range. If omitted, defaults to Unix epoch.
+
+    :type end: :class:`datetime.datetime`
+    :param end: (Optional) The (exclusive) upper bound of the timestamp
+                range. If omitted, no upper bound is used.
+    """
+
+    def __init__(self, start=None, end=None):
+        self.start = start
+        self.end = end
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return (other.start == self.start and
+                other.end == self.end)
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+    def to_pb(self):
+        """Converts the :class:`TimestampRange` to a protobuf.
+
+        :rtype: :class:`.data_pb2.TimestampRange`
+        :returns: The converted current object.
+        """
+        timestamp_range_kwargs = {}
+        if self.start is not None:
+            timestamp_range_kwargs['start_timestamp_micros'] = (
+                _microseconds_from_datetime(self.start))
+        if self.end is not None:
+            timestamp_range_kwargs['end_timestamp_micros'] = (
+                _microseconds_from_datetime(self.end))
+        return data_pb2.TimestampRange(**timestamp_range_kwargs)
+
+
+class TimestampRangeFilter(RowFilter):
+    """Row filter that limits cells to a range of time.
+
+    :type range_: :class:`TimestampRange`
+    :param range_: Range of time that cells should match against.
+    """
+
+    def __init__(self, range_):
+        self.range_ = range_
+
+    def __eq__(self, other):
+        if not isinstance(other, self.__class__):
+            return False
+        return other.range_ == self.range_
+
+    def to_pb(self):
+        """Converts the row filter to a protobuf.
+
+        First converts the ``range_`` on the current object to a protobuf and
+        then uses it in the ``timestamp_range_filter`` field.
+
+        :rtype: :class:`.data_pb2.RowFilter`
+        :returns: The converted current object.
+        """
+        return data_pb2.RowFilter(timestamp_range_filter=self.range_.to_pb())
+
+
 class ColumnRangeFilter(RowFilter):
     """A row filter to restrict to a range of columns.
 

gcloud/bigtable/test_row.py

@@ -398,6 +398,120 @@ def test_to_pb_exclusive_end(self):
         self.assertEqual(row_filter.to_pb(), expected_pb)
 
 
+class TestTimestampRange(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigtable.row import TimestampRange
+        return TimestampRange
+
+    def _makeOne(self, *args, **kwargs):
+        return self._getTargetClass()(*args, **kwargs)
+
+    def test_constructor(self):
+        start = object()
+        end = object()
+        time_range = self._makeOne(start=start, end=end)
+        self.assertTrue(time_range.start is start)
+        self.assertTrue(time_range.end is end)
+
+    def test___eq__(self):
+        start = object()
+        end = object()
+        time_range1 = self._makeOne(start=start, end=end)
+        time_range2 = self._makeOne(start=start, end=end)
+        self.assertEqual(time_range1, time_range2)
+
+    def test___eq__type_differ(self):
+        start = object()
+        end = object()
+        time_range1 = self._makeOne(start=start, end=end)
+        time_range2 = object()
+        self.assertNotEqual(time_range1, time_range2)
+
+    def test___ne__same_value(self):
+        start = object()
+        end = object()
+        time_range1 = self._makeOne(start=start, end=end)
+        time_range2 = self._makeOne(start=start, end=end)
+        comparison_val = (time_range1 != time_range2)
+        self.assertFalse(comparison_val)
+
+    def _to_pb_helper(self, start_micros=None, end_micros=None):
+        import datetime
+        from gcloud._helpers import _EPOCH
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+
+        pb_kwargs = {}
+
+        start = None
+        if start_micros is not None:
+            start = _EPOCH + datetime.timedelta(microseconds=start_micros)
+            pb_kwargs['start_timestamp_micros'] = start_micros
+        end = None
+        if end_micros is not None:
+            end = _EPOCH + datetime.timedelta(microseconds=end_micros)
+            pb_kwargs['end_timestamp_micros'] = end_micros
+        time_range = self._makeOne(start=start, end=end)
+
+        expected_pb = data_pb2.TimestampRange(**pb_kwargs)
+        self.assertEqual(time_range.to_pb(), expected_pb)
+
+    def test_to_pb(self):
+        # Makes sure already milliseconds granularity
+        start_micros = 30871000
+        end_micros = 12939371000
+        self._to_pb_helper(start_micros=start_micros,
+                           end_micros=end_micros)
+
+    def test_to_pb_start_only(self):
+        # Makes sure already milliseconds granularity
+        start_micros = 30871000
+        self._to_pb_helper(start_micros=start_micros)
+
+    def test_to_pb_end_only(self):
+        # Makes sure already milliseconds granularity
+        end_micros = 12939371000
+        self._to_pb_helper(end_micros=end_micros)
+
+
+class TestTimestampRangeFilter(unittest2.TestCase):
+
+    def _getTargetClass(self):
+        from gcloud.bigtable.row import TimestampRangeFilter
+        return TimestampRangeFilter
+
+    def _makeOne(self, *args, **kwargs):
+        return self._getTargetClass()(*args, **kwargs)
+
+    def test_constructor(self):
+        range_ = object()
+        row_filter = self._makeOne(range_)
+        self.assertTrue(row_filter.range_ is range_)
+
+    def test___eq__type_differ(self):
+        range_ = object()
+        row_filter1 = self._makeOne(range_)
+        row_filter2 = object()
+        self.assertNotEqual(row_filter1, row_filter2)
+
+    def test___eq__same_value(self):
+        range_ = object()
+        row_filter1 = self._makeOne(range_)
+        row_filter2 = self._makeOne(range_)
+        self.assertEqual(row_filter1, row_filter2)
+
+    def test_to_pb(self):
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+        from gcloud.bigtable.row import TimestampRange
+
+        range_ = TimestampRange()
+        row_filter = self._makeOne(range_)
+        pb_val = row_filter.to_pb()
+        expected_pb = data_pb2.RowFilter(
+            timestamp_range_filter=data_pb2.TimestampRange())
+        self.assertEqual(pb_val, expected_pb)
+
+
 class TestValueRegexFilter(unittest2.TestCase):
 
     def _getTargetClass(self):