Revert get_updates_for_prefix - reduce the amount of new code

Author: gwaramadze

quixstreams/state/base/transaction.py

@@ -6,6 +6,7 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Dict,
     Generic,
     Literal,
     Optional,
@@ -139,29 +140,15 @@ def get_column_families(self) -> Set[str]:
         """
         return set(self._updated.keys()) | set(self._deleted.keys())
 
-    def get_updates(self, cf_name: str = "default") -> dict[bytes, dict[bytes, bytes]]:
+    def get_updates(self, cf_name: str = "default") -> Dict[bytes, Dict[bytes, bytes]]:
         """
         Get all updated keys (excluding deleted)
-        in the format "{<prefix>: {<key>: <value>, ...}, ...}".
+        in the format "{<prefix>: {<key>: <value>}}".
 
         :param: cf_name: column family name
         """
         return self._updated.get(cf_name, {})
 
-    def get_updates_for_prefix(
-        self,
-        prefix: bytes,
-        cf_name: str = "default",
-    ) -> dict[bytes, bytes]:
-        """
-        Get all updated keys (excluding deleted)
-        in the format "{<key>: <value>, ...}".
-
-        :param: prefix: key prefix
-        :param: cf_name: column family name
-        """
-        return self._updated.get(cf_name, {}).get(prefix, {})
-
     def get_deletes(self, cf_name: str = "default") -> Set[bytes]:
         """
         Get all deleted keys (excluding updated) as a set.

quixstreams/state/rocksdb/timestamped.py

@@ -99,10 +99,7 @@ def get_last(
         value: Optional[bytes] = None
 
         deletes = self._update_cache.get_deletes(cf_name=cf_name)
-        updates = self._update_cache.get_updates_for_prefix(
-            prefix=prefix,
-            cf_name=cf_name,
-        )
+        updates = self._update_cache.get_updates(cf_name=cf_name).get(prefix, {})
 
         cached = sorted(updates.items(), reverse=True)
         for cached_key, cached_value in cached:
@@ -181,9 +178,7 @@ def _expire(
 
         key = self._serialize_key(max(timestamp - retention_ms, 0), prefix)
 
-        cached = self._update_cache.get_updates_for_prefix(
-            prefix=prefix, cf_name=cf_name
-        )
+        cached = self._update_cache.get_updates(cf_name=cf_name).get(prefix, {})
         # Cast to list to avoid RuntimeError: dictionary changed size during iteration
         for cached_key in list(cached):
             if cached_key < key:

tests/test_quixstreams/test_state/test_transaction.py

@@ -581,29 +581,6 @@ def test_get_updates_after_delete(self, cache: PartitionTransactionCache):
         cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
         assert cache.get_updates(cf_name="cf_name") == {b"prefix": {}}
 
-    def test_get_updates_for_prefix_empty(self, cache: PartitionTransactionCache):
-        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
-
-        # Delete an item and make sure it's not in "updates"
-        cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
-        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
-
-    def test_get_updates_for_prefix_present(self, cache: PartitionTransactionCache):
-        cache.set(key=b"key", value=b"value", prefix=b"prefix", cf_name="cf_name")
-        cache.set(key=b"key", value=b"value", prefix=b"other_prefix", cf_name="cf_name")
-        cache.set(key=b"key", value=b"value", prefix=b"other", cf_name="other_cf_name")
-
-        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {
-            b"key": b"value"
-        }
-
-    def test_get_updates_for_prefix_after_delete(
-        self, cache: PartitionTransactionCache
-    ):
-        cache.set(key=b"key", value=b"value", prefix=b"prefix", cf_name="cf_name")
-        cache.delete(key=b"key", prefix=b"prefix", cf_name="cf_name")
-        assert cache.get_updates_for_prefix(prefix=b"prefix", cf_name="cf_name") == {}
-
     def test_get_deletes_empty(self, cache: PartitionTransactionCache):
         assert cache.get_deletes(cf_name="cf_name") == set()