Lazy column family create (#905)

Author: gwaramadze

quixstreams/state/memory/partition.py

@@ -4,7 +4,6 @@
 
 from quixstreams.state import PartitionTransaction
 from quixstreams.state.base import PartitionTransactionCache, StorePartition
-from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.metadata import METADATA_CF_NAME, Marker
 from quixstreams.state.recovery import ChangelogProducer
 from quixstreams.utils.json import dumps as json_dumps
@@ -99,9 +98,6 @@ def write(
     def recover_from_changelog_message(
         self, key: bytes, value: Optional[bytes], cf_name: str, offset: int
     ) -> None:
-        if cf_name not in self._state:
-            raise ColumnFamilyDoesNotExist(f'Column family "{cf_name}" does not exist')
-
         if value:
             self._state.setdefault(cf_name, {})[key] = value
         else:

quixstreams/state/rocksdb/exceptions.py

@@ -1,9 +1,6 @@
 from quixstreams.state.exceptions import StateError
 
-__all__ = ("ColumnFamilyAlreadyExists",)
-
-
-class ColumnFamilyAlreadyExists(StateError): ...
+__all__ = ("RocksDBCorruptedError",)
 
 
 class RocksDBCorruptedError(StateError): ...

quixstreams/state/rocksdb/partition.py

@@ -17,12 +17,11 @@
     PartitionTransactionCache,
     StorePartition,
 )
-from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.metadata import METADATA_CF_NAME, Marker
 from quixstreams.state.recovery import ChangelogProducer
 from quixstreams.state.serialization import int_from_bytes, int_to_bytes
 
-from .exceptions import ColumnFamilyAlreadyExists, RocksDBCorruptedError
+from .exceptions import RocksDBCorruptedError
 from .metadata import (
     CHANGELOG_OFFSET_KEY,
 )
@@ -52,8 +51,6 @@ class RocksDBStorePartition(StorePartition):
     :param options: RocksDB options. If `None`, the default options will be used.
     """
 
-    additional_column_families: tuple[str, ...] = ()
-
     def __init__(
         self,
         path: str,
@@ -72,8 +69,6 @@ def __init__(
         self._db = self._init_rocksdb()
         self._cf_cache: Dict[str, Rdict] = {}
         self._cf_handle_cache: Dict[str, ColumnFamily] = {}
-        for cf_name in self.additional_column_families:
-            self._ensure_column_family(cf_name)
 
     def recover_from_changelog_message(
         self, key: bytes, value: Optional[bytes], cf_name: str, offset: int
@@ -149,7 +144,9 @@ def get(
         :param cf_name: rocksdb column family name. Default - "default"
         :return: a value if the key is present in the DB. Otherwise, `default`
         """
-        result = self.get_column_family(cf_name).get(key, default=Marker.UNDEFINED)
+        result = self.get_or_create_column_family(cf_name).get(
+            key, default=Marker.UNDEFINED
+        )
 
         # RDict accept Any type as value but we only write bytes so we should only get bytes back.
         return cast(Union[bytes, Literal[Marker.UNDEFINED]], result)
@@ -172,7 +169,7 @@ def iter_items(
             Default is "default".
         :return: An iterator yielding (key, value) tuples.
         """
-        cf = self.get_column_family(cf_name=cf_name)
+        cf = self.get_or_create_column_family(cf_name=cf_name)
 
         # Set iterator bounds to reduce IO by limiting the range of keys fetched
         read_opt = ReadOptions()
@@ -219,15 +216,15 @@ def exists(self, key: bytes, cf_name: str = "default") -> bool:
         :param cf_name: rocksdb column family name. Default - "default"
         :return: `True` if the key is present, `False` otherwise.
         """
-        cf_dict = self.get_column_family(cf_name)
+        cf_dict = self.get_or_create_column_family(cf_name)
         return key in cf_dict
 
     def get_changelog_offset(self) -> Optional[int]:
         """
         Get offset that the changelog is up-to-date with.
         :return: offset or `None` if there's no processed offset yet
         """
-        metadata_cf = self.get_column_family(METADATA_CF_NAME)
+        metadata_cf = self.get_or_create_column_family(METADATA_CF_NAME)
         offset_bytes = metadata_cf.get(CHANGELOG_OFFSET_KEY)
         if offset_bytes is None:
             return None
@@ -288,18 +285,12 @@ def get_column_family_handle(self, cf_name: str) -> ColumnFamily:
         :return: instance of `rocksdict.ColumnFamily`
         """
         if (cf_handle := self._cf_handle_cache.get(cf_name)) is None:
-            try:
-                cf_handle = self._db.get_column_family_handle(cf_name)
-                self._cf_handle_cache[cf_name] = cf_handle
-            except Exception as exc:
-                if "does not exist" in str(exc):
-                    raise ColumnFamilyDoesNotExist(
-                        f'Column family "{cf_name}" does not exist'
-                    )
-                raise
+            self.get_or_create_column_family(cf_name)
+            cf_handle = self._db.get_column_family_handle(cf_name)
+            self._cf_handle_cache[cf_name] = cf_handle
         return cf_handle
 
-    def get_column_family(self, cf_name: str) -> Rdict:
+    def get_or_create_column_family(self, cf_name: str) -> Rdict:
         """
         Get a column family instance.
         This method will cache the CF instance to avoid creating them repeatedly.
@@ -310,38 +301,14 @@ def get_column_family(self, cf_name: str) -> Rdict:
         if (cf := self._cf_cache.get(cf_name)) is None:
             try:
                 cf = self._db.get_column_family(cf_name)
-                self._cf_cache[cf_name] = cf
             except Exception as exc:
-                if "does not exist" in str(exc):
-                    raise ColumnFamilyDoesNotExist(
-                        f'Column family "{cf_name}" does not exist'
-                    )
-                raise
-        return cf
-
-    def create_column_family(self, cf_name: str):
-        try:
-            cf = self._db.create_column_family(cf_name, options=self._rocksdb_options)
-        except Exception as exc:
-            if "column family already exists" in str(exc).lower():
-                raise ColumnFamilyAlreadyExists(
-                    f'Column family already exists: "{cf_name}"'
-                )
-            raise
-
-        self._cf_cache[cf_name] = cf
-
-    def drop_column_family(self, cf_name: str):
-        self._cf_cache.pop(cf_name, None)
-        self._cf_handle_cache.pop(cf_name, None)
-        try:
-            self._db.drop_column_family(cf_name)
-        except Exception as exc:
-            if "invalid column family:" in str(exc).lower():
-                raise ColumnFamilyDoesNotExist(
-                    f'Column family does not exist: "{cf_name}"'
+                if "does not exist" not in str(exc):
+                    raise
+                cf = self._db.create_column_family(
+                    cf_name, options=self._rocksdb_options
                 )
-            raise
+            self._cf_cache[cf_name] = cf
+        return cf
 
     def list_column_families(self) -> List[str]:
         return self._db.list_cf(self._path)
@@ -427,9 +394,3 @@ def _update_changelog_offset(self, batch: WriteBatch, offset: int):
             int_to_bytes(offset),
             self.get_column_family_handle(METADATA_CF_NAME),
         )
-
-    def _ensure_column_family(self, cf_name: str) -> None:
-        try:
-            self.get_column_family(cf_name)
-        except ColumnFamilyDoesNotExist:
-            self.create_column_family(cf_name)

quixstreams/state/rocksdb/timestamped.py

@@ -262,8 +262,6 @@ class TimestampedStorePartition(RocksDBStorePartition):
     `TimestampedPartitionTransaction` instances to handle atomic operations for that partition.
     """
 
-    additional_column_families = (MIN_ELIGIBLE_TIMESTAMPS_CF_NAME,)
-
     def begin(self) -> TimestampedPartitionTransaction:
         # Override the method to specify the correct return type
         return TimestampedPartitionTransaction(

quixstreams/state/rocksdb/windowed/partition.py

@@ -2,14 +2,6 @@
 from typing import Iterator, cast
 
 from ..partition import RocksDBStorePartition
-from .metadata import (
-    GLOBAL_COUNTER_CF_NAME,
-    LATEST_DELETED_VALUE_CF_NAME,
-    LATEST_DELETED_WINDOW_CF_NAME,
-    LATEST_EXPIRED_WINDOW_CF_NAME,
-    LATEST_TIMESTAMPS_CF_NAME,
-    VALUES_CF_NAME,
-)
 from .transaction import WindowedRocksDBPartitionTransaction
 
 logger = logging.getLogger(__name__)
@@ -24,15 +16,6 @@ class WindowedRocksDBStorePartition(RocksDBStorePartition):
     stores the expiration index to delete expired windows.
     """
 
-    additional_column_families = (
-        LATEST_DELETED_VALUE_CF_NAME,
-        LATEST_EXPIRED_WINDOW_CF_NAME,
-        LATEST_DELETED_WINDOW_CF_NAME,
-        LATEST_TIMESTAMPS_CF_NAME,
-        GLOBAL_COUNTER_CF_NAME,
-        VALUES_CF_NAME,
-    )
-
     def iter_keys(self, cf_name: str = "default") -> Iterator[bytes]:
         """
         Iterate over all keys in the DB.
@@ -42,7 +25,7 @@ def iter_keys(self, cf_name: str = "default") -> Iterator[bytes]:
         :param cf_name: rocksdb column family name. Default - "default"
         :return: An iterable of keys
         """
-        cf_dict = self.get_column_family(cf_name)
+        cf_dict = self.get_or_create_column_family(cf_name)
         return cast(Iterator[bytes], cf_dict.keys())
 
     def begin(self) -> WindowedRocksDBPartitionTransaction:

tests/test_quixstreams/test_state/test_partition.py

@@ -1,7 +1,7 @@
 import pytest
 
-from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.manager import SUPPORTED_STORES
+from quixstreams.state.metadata import Marker
 
 
 @pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
@@ -24,13 +24,12 @@ def test_recover_from_changelog_message_success(self, store_value, store_partiti
         assert store_partition.get(b"key") == b"value"
         assert store_partition.get_changelog_offset() == 1
 
-    def test_recover_from_changelog_message_missing_cf_fails(self, store_partition):
-        with pytest.raises(
-            ColumnFamilyDoesNotExist, match='Column family "some_cf" does not exist'
-        ):
-            store_partition.recover_from_changelog_message(
-                key=b"key", value=b"value", cf_name="some_cf", offset=1
-            )
+    def test_recover_from_changelog_message_missing_cf(self, store_partition):
+        store_partition.recover_from_changelog_message(
+            key=b"key", value=b"value", cf_name="some_cf", offset=1
+        )
+        assert store_partition.get(b"key") == Marker.UNDEFINED
+        assert store_partition.get_changelog_offset() == 1
 
     def test_write_changelog_offset(self, store_partition):
         assert store_partition.get_changelog_offset() is None

tests/test_quixstreams/test_state/test_rocksdb/test_rocksdb_partition.py

@@ -5,9 +5,7 @@
 import pytest
 from rocksdict import Rdict
 
-from quixstreams.state.exceptions import ColumnFamilyDoesNotExist
 from quixstreams.state.rocksdb import (
-    ColumnFamilyAlreadyExists,
     RocksDBOptions,
     RocksDBStorePartition,
 )
@@ -113,47 +111,19 @@ def test_db_corrupted_sst_file(self, store_partition_factory, tmp_path):
 
         store_partition_factory(options=RocksDBOptions(on_corrupted_recreate=True))
 
-    def test_create_and_get_column_family(self, store_partition: RocksDBStorePartition):
-        store_partition.create_column_family("cf")
-        assert store_partition.get_column_family("cf")
+    def test_get_or_create_column_family(self, store_partition: RocksDBStorePartition):
+        assert store_partition.get_or_create_column_family("cf")
 
-    def test_create_column_family_already_exists(
+    def test_get_or_create_column_family_cached(
         self, store_partition: RocksDBStorePartition
     ):
-        store_partition.create_column_family("cf")
-        with pytest.raises(ColumnFamilyAlreadyExists):
-            store_partition.create_column_family("cf")
-
-    def test_get_column_family_doesnt_exist(
-        self, store_partition: RocksDBStorePartition
-    ):
-        with pytest.raises(ColumnFamilyDoesNotExist):
-            store_partition.get_column_family("cf")
-
-    def test_get_column_family_cached(self, store_partition: RocksDBStorePartition):
-        store_partition.create_column_family("cf")
-        cf1 = store_partition.get_column_family("cf")
-        cf2 = store_partition.get_column_family("cf")
+        cf1 = store_partition.get_or_create_column_family("cf")
+        cf2 = store_partition.get_or_create_column_family("cf")
         assert cf1 is cf2
 
-    def test_create_and_drop_column_family(
-        self, store_partition: RocksDBStorePartition
-    ):
-        store_partition.create_column_family("cf")
-        store_partition.drop_column_family("cf")
-
-        with pytest.raises(ColumnFamilyDoesNotExist):
-            store_partition.get_column_family("cf")
-
-    def test_drop_column_family_doesnt_exist(
-        self, store_partition: RocksDBStorePartition
-    ):
-        with pytest.raises(ColumnFamilyDoesNotExist):
-            store_partition.drop_column_family("cf")
-
     def test_list_column_families(self, store_partition: RocksDBStorePartition):
-        store_partition.create_column_family("cf1")
-        store_partition.create_column_family("cf2")
+        store_partition.get_or_create_column_family("cf1")
+        store_partition.get_or_create_column_family("cf2")
         cfs = store_partition.list_column_families()
         assert "cf1" in cfs
         assert "cf2" in cfs
@@ -187,7 +157,7 @@ def test_list_column_families_defaults(
         ]
 
     def test_ensure_metadata_cf(self, store_partition: RocksDBStorePartition):
-        assert store_partition.get_column_family("__metadata__")
+        assert store_partition.get_or_create_column_family("__metadata__")
 
     @pytest.mark.parametrize(
         ["backwards", "expected"],

tests/test_quixstreams/test_state/test_rocksdb/test_rocksdb_transaction.py

@@ -8,7 +8,7 @@ def test_set_get_with_column_family(self, store_partition):
         key = "key"
         value = "value"
         prefix = b"__key__"
-        store_partition.create_column_family("cf")
+        store_partition.get_or_create_column_family("cf")
 
         with store_partition.begin() as tx:
             tx.set(key, value, cf_name="cf", prefix=prefix)
@@ -21,7 +21,7 @@ def test_set_delete_get_with_column_family(self, store_partition):
         key = "key"
         value = "value"
         prefix = b"__key__"
-        store_partition.create_column_family("cf")
+        store_partition.get_or_create_column_family("cf")
 
         with store_partition.begin() as tx:
             tx.set(key, value, cf_name="cf", prefix=prefix)
@@ -35,7 +35,7 @@ def test_set_delete_get_with_column_family(self, store_partition):
     def test_set_exists_get_with_column_family(self, store_partition):
         key = "key"
         value = "value"
-        store_partition.create_column_family("cf")
+        store_partition.get_or_create_column_family("cf")
         prefix = b"__key__"
 
         with store_partition.begin() as tx: