Add retention

Author: gwaramadze

quixstreams/state/rocksdb/timestamped.py

@@ -17,6 +17,8 @@
     "TimestampedPartitionTransaction",
 )
 
+DAYS_7 = 7 * 24 * 60 * 60 * 1000
+
 
 class TimestampedPartitionTransaction(PartitionTransaction):
     """
@@ -39,6 +41,7 @@ def get_last(
         self,
         timestamp: int,
         prefix: Any,
+        retention: int = DAYS_7,
         cf_name: str = "default",
     ) -> Optional[Any]:
         """Get the latest value for a prefix up to a given timestamp.
@@ -56,8 +59,12 @@ def get_last(
         :return: The deserialized value if found, otherwise None.
         """
         prefix = self._ensure_bytes(prefix)
-        # Add +1 because the storage `.iter_items()` is exclusive on the upper bound
-        key = self._serialize_key(timestamp + 1, prefix)
+
+        # Negative retention is not allowed
+        lower_bound = self._serialize_key(max(timestamp - retention, 0), prefix)
+        # +1 because upper bound is exclusive
+        upper_bound = self._serialize_key(timestamp + 1, prefix)
+
         value: Optional[bytes] = None
 
         deletes = self._update_cache.get_deletes(cf_name=cf_name)
@@ -68,13 +75,13 @@ def get_last(
 
         cached = sorted(updates.items(), reverse=True)
         for cached_key, cached_value in cached:
-            if prefix < cached_key < key and cached_key not in deletes:
+            if lower_bound <= cached_key < upper_bound and cached_key not in deletes:
                 value = cached_value
                 break
 
         stored = self._partition.iter_items(
-            lower_bound=prefix,
-            upper_bound=key,
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
             backwards=True,
             cf_name=cf_name,
         )
@@ -108,7 +115,19 @@ def set(self, timestamp: int, value: Any, prefix: Any, cf_name: str = "default")
         super().set(timestamp, value, prefix, cf_name=cf_name)
 
     def expire(self, timestamp: int, prefix: bytes, cf_name: str = "default"):
-        key = self._serialize_key(timestamp + 1, prefix)
+        """
+        Delete all entries for a given prefix with timestamps less than the
+        provided timestamp.
+
+        This applies to both the in-memory update cache and the underlying
+        RocksDB store within the current transaction.
+
+        :param timestamp: The upper bound timestamp (exclusive) in milliseconds.
+            Entries with timestamps strictly less than this will be deleted.
+        :param prefix: The key prefix.
+        :param cf_name: Column family name.
+        """
+        key = self._serialize_key(timestamp, prefix)
 
         cached = self._update_cache.get_updates_for_prefix(
             prefix=prefix,

tests/test_quixstreams/test_state/test_rocksdb/test_timestamped.py

@@ -54,7 +54,7 @@ def test_get_last_ignore_deleted(
         tx.set(timestamp=9, value="value9-stored", prefix=b"key")
 
     with transaction() as tx:
-        tx.expire(timestamp=9, prefix=b"key")
+        tx.expire(timestamp=10, prefix=b"key")
 
         # Message with timestamp 8 comes out of order after later messages
         # got expired in the same transaction.
@@ -114,13 +114,35 @@ def test_get_last_prefix_not_bytes(transaction: TimestampedPartitionTransaction)
         assert tx.get_last(timestamp=10, prefix=b'"key"') == "value"
 
 
+def test_get_last_from_cache_with_retention(
+    transaction: TimestampedPartitionTransaction,
+):
+    with transaction() as tx:
+        tx.set(timestamp=5, value="value", prefix=b"key")
+        assert tx.get_last(timestamp=10, prefix=b"key") == "value"
+        assert tx.get_last(timestamp=10, prefix=b"key", retention=5) == "value"
+        assert tx.get_last(timestamp=10, prefix=b"key", retention=4) == None
+
+
+def test_get_last_from_store_with_retention(
+    transaction: TimestampedPartitionTransaction,
+):
+    with transaction() as tx:
+        tx.set(timestamp=5, value="value", prefix=b"key")
+
+    with transaction() as tx:
+        assert tx.get_last(timestamp=10, prefix=b"key") == "value"
+        assert tx.get_last(timestamp=10, prefix=b"key", retention=5) == "value"
+        assert tx.get_last(timestamp=10, prefix=b"key", retention=4) == None
+
+
 def test_expire_cached(transaction: TimestampedPartitionTransaction):
     with transaction() as tx:
         tx.set(timestamp=1, value="value1", prefix=b"key")
         tx.set(timestamp=10, value="value10", prefix=b"key")
         tx.set(timestamp=11, value="value11", prefix=b"key")
 
-        tx.expire(timestamp=10, prefix=b"key")
+        tx.expire(timestamp=11, prefix=b"key")
 
         assert tx.get_last(timestamp=10, prefix=b"key") == None
         assert tx.get_last(timestamp=11, prefix=b"key") == "value11"
@@ -133,7 +155,7 @@ def test_expire_stored(transaction: TimestampedPartitionTransaction):
         tx.set(timestamp=11, value="value11", prefix=b"key")
 
     with transaction() as tx:
-        tx.expire(timestamp=10, prefix=b"key")
+        tx.expire(timestamp=11, prefix=b"key")
 
         assert tx.get_last(timestamp=10, prefix=b"key") == None
         assert tx.get_last(timestamp=11, prefix=b"key") == "value11"
@@ -146,7 +168,7 @@ def test_expire_idempotent(transaction: TimestampedPartitionTransaction):
     with transaction() as tx:
         tx.set(timestamp=10, value="value10", prefix=b"key")
 
-        tx.expire(timestamp=10, prefix=b"key")
-        tx.expire(timestamp=10, prefix=b"key")
+        tx.expire(timestamp=11, prefix=b"key")
+        tx.expire(timestamp=11, prefix=b"key")
 
         assert tx.get_last(timestamp=10, prefix=b"key") == None