Ensure additional column families in RocksDBStorePartition base class

Author: gwaramadze

quixstreams/state/rocksdb/partition.py

@@ -42,6 +42,8 @@ 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,
@@ -60,6 +62,8 @@ 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

quixstreams/state/rocksdb/timestamped.py

@@ -12,7 +12,6 @@
     LATEST_TIMESTAMP_KEY,
     LATEST_TIMESTAMPS_CF_NAME,
 )
-from quixstreams.state.rocksdb.types import RocksDBOptionsType
 from quixstreams.state.serialization import (
     DumpsFunc,
     LoadsFunc,
@@ -234,15 +233,7 @@ class TimestampedStorePartition(RocksDBStorePartition):
     """
 
     partition_transaction_class = TimestampedPartitionTransaction
-
-    def __init__(
-        self,
-        path: str,
-        options: Optional[RocksDBOptionsType] = None,
-        changelog_producer: Optional[ChangelogProducer] = None,
-    ) -> None:
-        super().__init__(path, options=options, changelog_producer=changelog_producer)
-        self._ensure_column_family(LATEST_TIMESTAMPS_CF_NAME)
+    additional_column_families = (LATEST_TIMESTAMPS_CF_NAME,)
 
 
 class TimestampedStore(RocksDBStore):

quixstreams/state/rocksdb/windowed/partition.py

@@ -1,11 +1,8 @@
 import logging
-from typing import Iterator, Optional, cast
-
-from quixstreams.state.recovery import ChangelogProducer
+from typing import Iterator, cast
 
 from ..metadata import LATEST_TIMESTAMPS_CF_NAME
 from ..partition import RocksDBStorePartition
-from ..types import RocksDBOptionsType
 from .metadata import (
     GLOBAL_COUNTER_CF_NAME,
     LATEST_DELETED_VALUE_CF_NAME,
@@ -25,28 +22,17 @@ class WindowedRocksDBStorePartition(RocksDBStorePartition):
 
     Besides the data, it keeps track of the latest observed timestamp and
     stores the expiration index to delete expired windows.
-
-    :param path: an absolute path to the RocksDB folder
-    :param options: RocksDB options. If `None`, the default options will be used.
     """
 
     partition_transaction_class = WindowedRocksDBPartitionTransaction
-
-    def __init__(
-        self,
-        path: str,
-        options: Optional[RocksDBOptionsType] = None,
-        changelog_producer: Optional[ChangelogProducer] = None,
-    ):
-        super().__init__(
-            path=path, options=options, changelog_producer=changelog_producer
-        )
-        self._ensure_column_family(LATEST_DELETED_VALUE_CF_NAME)
-        self._ensure_column_family(LATEST_EXPIRED_WINDOW_CF_NAME)
-        self._ensure_column_family(LATEST_DELETED_WINDOW_CF_NAME)
-        self._ensure_column_family(LATEST_TIMESTAMPS_CF_NAME)
-        self._ensure_column_family(GLOBAL_COUNTER_CF_NAME)
-        self._ensure_column_family(VALUES_CF_NAME)
+    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]:
         """