Refactor

Author: gwaramadze

quixstreams/state/rocksdb/timestamped.py

@@ -65,7 +65,6 @@ def get_last(
         self,
         timestamp: int,
         prefix: Any,
-        retention_ms: int = DAYS_7,
         cf_name: str = "default",
     ) -> Optional[Any]:
         """Get the latest value for a prefix up to a given timestamp.
@@ -79,15 +78,12 @@ def get_last(
 
         :param timestamp: The upper bound timestamp (inclusive) in milliseconds.
         :param prefix: The key prefix.
-        :param retention_ms: The retention period in milliseconds.
         :param cf_name: The column family name.
         :return: The deserialized value if found, otherwise None.
         """
 
         prefix = self._ensure_bytes(prefix)
-        min_eligible_timestamp = self._get_min_eligible_timestamp(
-            prefix, timestamp, retention_ms
-        )
+        min_eligible_timestamp = self._get_min_eligible_timestamp(prefix)
 
         lower_bound = self._serialize_key(min_eligible_timestamp, prefix)
         # +1 because upper bound is exclusive
@@ -147,33 +143,26 @@ def set_for_timestamp(
         """
         prefix = self._ensure_bytes(prefix)
         super().set(timestamp, value, prefix, cf_name=cf_name)
-        self._expire(
-            timestamp=timestamp,
-            prefix=prefix,
-            retention_ms=retention_ms,
-            cf_name=cf_name,
+        min_eligible_timestamp = max(
+            self._get_min_eligible_timestamp(prefix),
+            timestamp - retention_ms,
         )
+        self._set_min_eligible_timestamp(prefix, min_eligible_timestamp)
+        self._expire(prefix=prefix, cf_name=cf_name)
 
-    def _expire(
-        self, timestamp: int, prefix: bytes, retention_ms: int, cf_name: str = "default"
-    ) -> None:
+    def _expire(self, prefix: bytes, cf_name: str = "default") -> None:
         """
         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.
         """
 
-        min_eligible_timestamp = self._get_min_eligible_timestamp(
-            prefix, timestamp, retention_ms
-        )
-        self._set_min_eligible_timestamp(prefix, min_eligible_timestamp)
+        min_eligible_timestamp = self._get_min_eligible_timestamp(prefix)
 
         key = self._serialize_key(min_eligible_timestamp, prefix)
 
@@ -203,16 +192,13 @@ def _serialize_key(self, key: Union[int, bytes], prefix: bytes) -> bytes:
             return prefix + SEPARATOR + key
         raise ValueError(f"Invalid key type: {type(key)}")
 
-    def _get_min_eligible_timestamp(
-        self, prefix: bytes, timestamp: int, retention_ms: int
-    ) -> Any:
+    def _get_min_eligible_timestamp(self, prefix: bytes) -> int:
         cache = self._min_eligible_timestamps
-        ts = (
+        return (
             cache.timestamps.get(prefix)
             or self.get(key=cache.key, prefix=prefix, cf_name=cache.cf_name)
             or 0
         )
-        return max(ts, timestamp - retention_ms)
 
     def _set_min_eligible_timestamp(self, prefix: bytes, timestamp: int) -> None:
         cache = self._min_eligible_timestamps

tests/test_quixstreams/test_state/test_rocksdb/test_timestamped.py

@@ -101,44 +101,22 @@ 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_for_timestamp(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_ms=5) == "value"
-        assert tx.get_last(timestamp=10, prefix=b"key", retention_ms=4) == None
-
-
-def test_get_last_from_store_with_retention(
-    transaction: TimestampedPartitionTransaction,
-):
-    with transaction() as tx:
-        tx.set_for_timestamp(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_ms=5) == "value"
-        assert tx.get_last(timestamp=10, prefix=b"key", retention_ms=4) == None
-
-
 def test_get_last_for_out_of_order_timestamp(
     transaction: TimestampedPartitionTransaction,
 ):
     with transaction() as tx:
         tx.set_for_timestamp(
             timestamp=10, value="value10", prefix=b"key", retention_ms=5
         )
-        assert tx.get_last(timestamp=10, prefix=b"key", retention_ms=5) == "value10"
+        assert tx.get_last(timestamp=10, prefix=b"key") == "value10"
         tx.set_for_timestamp(timestamp=5, value="value5", prefix=b"key", retention_ms=5)
         tx.set_for_timestamp(timestamp=4, value="value4", prefix=b"key", retention_ms=5)
 
     with transaction() as tx:
-        assert tx.get_last(timestamp=5, prefix=b"key", retention_ms=5) == "value5"
+        assert tx.get_last(timestamp=5, prefix=b"key") == "value5"
 
         # Retention watermark is 10 - 5 = 5 so everything lower is ignored
-        assert tx.get_last(timestamp=4, prefix=b"key", retention_ms=5) is None
+        assert tx.get_last(timestamp=4, prefix=b"key") is None
 
 
 def test_set_for_timestamp_with_prefix_not_bytes(
@@ -156,8 +134,8 @@ def test_set_for_timestamp_with_retention_cached(
     with transaction() as tx:
         tx.set_for_timestamp(timestamp=2, value="v2", prefix=b"key", retention_ms=2)
         tx.set_for_timestamp(timestamp=5, value="v5", prefix=b"key", retention_ms=2)
-        assert tx.get_last(timestamp=2, prefix=b"key", retention_ms=3) == None
-        assert tx.get_last(timestamp=5, prefix=b"key", retention_ms=3) == "v5"
+        assert tx.get_last(timestamp=2, prefix=b"key") == None
+        assert tx.get_last(timestamp=5, prefix=b"key") == "v5"
 
 
 def test_set_for_timestamp_with_retention_stored(
@@ -168,5 +146,5 @@ def test_set_for_timestamp_with_retention_stored(
         tx.set_for_timestamp(timestamp=5, value="v5", prefix=b"key", retention_ms=2)
 
     with transaction() as tx:
-        assert tx.get_last(timestamp=2, prefix=b"key", retention_ms=3) == None
-        assert tx.get_last(timestamp=5, prefix=b"key", retention_ms=3) == "v5"
+        assert tx.get_last(timestamp=2, prefix=b"key") == None
+        assert tx.get_last(timestamp=5, prefix=b"key") == "v5"