app: Add option to select store backend (#544)

Add support for multiple store backend

Author: ovv

quixstreams/state/manager.py

@@ -1,7 +1,7 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import List, Dict, Optional, Union
+from typing import List, Dict, Optional, Union, Type, get_args
 
 from quixstreams.rowproducer import RowProducer
 from .exceptions import (
@@ -14,12 +14,15 @@
 from .rocksdb.windowed.store import WindowedRocksDBStore
 from .base import Store, StorePartition
 
-__all__ = ("StateStoreManager", "DEFAULT_STATE_STORE_NAME")
+__all__ = ("StateStoreManager", "DEFAULT_STATE_STORE_NAME", "StoreTypes")
 
 logger = logging.getLogger(__name__)
 
 DEFAULT_STATE_STORE_NAME = "default"
 
+StoreTypes = Union[Type[RocksDBStore]]
+SUPPORTED_STORES = get_args(StoreTypes)
+
 
 class StateStoreManager:
     """
@@ -38,12 +41,14 @@ def __init__(
         rocksdb_options: Optional[RocksDBOptionsType] = None,
         producer: Optional[RowProducer] = None,
         recovery_manager: Optional[RecoveryManager] = None,
+        default_store_type: StoreTypes = RocksDBStore,
     ):
         self._state_dir = (Path(state_dir) / group_id).absolute()
         self._rocksdb_options = rocksdb_options
         self._stores: Dict[str, Dict[str, Store]] = {}
         self._producer = producer
         self._recovery_manager = recovery_manager
+        self._default_store_type = default_store_type
 
     def _init_state_dir(self):
         logger.info(f'Initializing state directory at "{self._state_dir}"')
@@ -84,6 +89,10 @@ def using_changelogs(self) -> bool:
         """
         return bool(self._recovery_manager)
 
+    @property
+    def default_store_type(self) -> StoreTypes:
+        return self._default_store_type
+
     def do_recovery(self):
         """
         Perform a state recovery, if necessary.
@@ -130,7 +139,10 @@ def _setup_changelogs(
             )
 
     def register_store(
-        self, topic_name: str, store_name: str = DEFAULT_STATE_STORE_NAME
+        self,
+        topic_name: str,
+        store_name: str = DEFAULT_STATE_STORE_NAME,
+        store_type: Optional[StoreTypes] = None,
     ):
         """
         Register a state store to be managed by StateStoreManager.
@@ -142,17 +154,25 @@ def register_store(
 
         :param topic_name: topic name
         :param store_name: store name
+        :param store_type: the storage type used for this store.
+            Default to StateStoreManager `default_store_type`
         """
         if self._stores.get(topic_name, {}).get(store_name) is None:
-            self._stores.setdefault(topic_name, {})[store_name] = RocksDBStore(
-                name=store_name,
-                topic=topic_name,
-                base_dir=str(self._state_dir),
-                changelog_producer_factory=self._setup_changelogs(
-                    topic_name, store_name
-                ),
-                options=self._rocksdb_options,
-            )
+            changelog_producer_factory = self._setup_changelogs(topic_name, store_name)
+
+            store_type = store_type or self.default_store_type
+            if store_type == RocksDBStore:
+                factory = RocksDBStore(
+                    name=store_name,
+                    topic=topic_name,
+                    base_dir=str(self._state_dir),
+                    changelog_producer_factory=changelog_producer_factory,
+                    options=self._rocksdb_options,
+                )
+            else:
+                raise ValueError(f"invalid store type: {store_type}")
+
+            self._stores.setdefault(topic_name, {})[store_name] = factory
 
     def register_windowed_store(self, topic_name: str, store_name: str):
         """

tests/conftest.py

@@ -15,7 +15,6 @@
     "tests.test_quixstreams.test_models.test_serializers.fixtures",
     "tests.test_quixstreams.test_platforms.test_quix.fixtures",
     "tests.test_quixstreams.test_state.fixtures",
-    "tests.test_quixstreams.test_state.test_rocksdb.fixtures",
     "tests.test_quixstreams.test_state.test_rocksdb.test_windowed.fixtures",
 ]
 

tests/test_quixstreams/fixtures.py

@@ -1,7 +1,7 @@
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from typing import Optional, Union
-from unittest.mock import create_autospec, patch
+from unittest.mock import create_autospec, patch, PropertyMock
 
 import pytest
 from confluent_kafka.admin import (
@@ -47,6 +47,7 @@
 from quixstreams.rowconsumer import RowConsumer
 from quixstreams.rowproducer import RowProducer
 from quixstreams.state import StateStoreManager
+from quixstreams.state.manager import StoreTypes
 from quixstreams.state.recovery import RecoveryManager
 
 
@@ -278,7 +279,7 @@ def factory(
 
 
 @pytest.fixture()
-def app_factory(kafka_container, random_consumer_group, tmp_path):
+def app_factory(kafka_container, random_consumer_group, tmp_path, store_type):
     def factory(
         consumer_group: Optional[str] = None,
         auto_offset_reset: AutoOffsetReset = "latest",
@@ -296,6 +297,7 @@ def factory(
         topic_manager: Optional[TopicManager] = None,
         processing_guarantee: ProcessingGuarantee = "at-least-once",
         request_timeout: float = 30,
+        store_type: StoreTypes = store_type,
     ) -> Application:
         state_dir = state_dir or (tmp_path / "state").absolute()
         return Application(
@@ -318,16 +320,22 @@ def factory(
             request_timeout=request_timeout,
         )
 
-    return factory
+    with patch(
+        "quixstreams.state.manager.StateStoreManager.default_store_type",
+        new_callable=PropertyMock,
+    ) as m:
+        m.return_value = store_type
+        yield factory
 
 
 @pytest.fixture()
-def state_manager_factory(tmp_path):
+def state_manager_factory(store_type, tmp_path):
     def factory(
         group_id: Optional[str] = None,
         state_dir: Optional[str] = None,
         producer: Optional[RowProducer] = None,
         recovery_manager: Optional[RecoveryManager] = None,
+        default_store_type: StoreTypes = store_type,
     ) -> StateStoreManager:
         group_id = group_id or str(uuid.uuid4())
         state_dir = state_dir or str(uuid.uuid4())
@@ -336,6 +344,7 @@ def factory(
             state_dir=str(tmp_path / state_dir),
             producer=producer,
             recovery_manager=recovery_manager,
+            default_store_type=default_store_type,
         )
 
     return factory
@@ -434,6 +443,7 @@ def quix_app_factory(
     topic_admin,
     quix_mock_config_builder_factory,
     quix_topic_manager_factory,
+    store_type,
 ):
     """
     For doing testing with Quix Applications against a local cluster.
@@ -454,6 +464,7 @@ def factory(
         auto_create_topics: bool = True,
         use_changelog_topics: bool = True,
         workspace_id: str = "my_ws",
+        store_type: Optional[StoreTypes] = store_type,
     ) -> Application:
         state_dir = state_dir or (tmp_path / "state").absolute()
         return Application(
@@ -474,7 +485,12 @@ def factory(
             ),
         )
 
-    return factory
+    with patch(
+        "quixstreams.state.manager.StateStoreManager.default_store_type",
+        new_callable=PropertyMock,
+    ) as m:
+        m.return_value = store_type
+        yield factory
 
 
 @pytest.fixture()

tests/test_quixstreams/test_app.py

@@ -30,6 +30,7 @@
 from quixstreams.rowproducer import RowProducer
 from quixstreams.sinks import SinkBatch, SinkBackpressureError
 from quixstreams.state import State
+from quixstreams.state.manager import SUPPORTED_STORES
 from quixstreams.sources import SourceException, multiprocessing
 from tests.utils import DummySink, DummySource
 
@@ -1036,6 +1037,7 @@ def test_topic_name_and_config(self, quix_app_factory):
         assert expected_topic.config.num_partitions == topic_partitions
 
 
+@pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
 class TestQuixApplicationWithState:
     def test_quix_app_no_state_management_warning(
         self, quix_app_factory, monkeypatch, topic_factory, executor
@@ -1086,6 +1088,7 @@ def test_quix_app_state_dir_mismatch_warning(
         assert "does not match the state directory" in warning
 
 
+@pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
 class TestApplicationWithState:
     def test_run_stateful_success(
         self,
@@ -1393,6 +1396,7 @@ def test_app_use_changelog_false(self, app_factory):
         assert not app._state_manager.using_changelogs
 
 
+@pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
 class TestApplicationRecovery:
     def test_changelog_recovery_default_store(
         self,
@@ -1672,6 +1676,7 @@ def validate_state():
     @pytest.mark.parametrize("processing_guarantee", ["at-least-once", "exactly-once"])
     def test_changelog_recovery_consistent_after_failed_commit(
         self,
+        store_type,
         app_factory,
         executor,
         tmp_path,
@@ -1734,6 +1739,7 @@ def get_app():
                 consumer_group=consumer_group,
                 state_dir=state_dir,
                 processing_guarantee=processing_guarantee,
+                store_type=store_type,
             )
             topic = app.topic(topic_name)
             sdf = app.dataframe(topic)
@@ -2376,6 +2382,7 @@ def on_message_processed(*_):
                     "groupby_timestamp": timestamp,
                 }
 
+    @pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
     def test_stateful(
         self,
         app_factory,
@@ -2466,6 +2473,7 @@ def count(_, state: State):
                 # All keys in state must be prefixed with the message key
                 assert tx.get("total", prefix=message_key) == messages_per_topic
 
+    @pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
     def test_changelog_recovery(
         self,
         app_factory,

tests/test_quixstreams/test_checkpointing.py

@@ -16,6 +16,7 @@
 from quixstreams.sinks import SinkManager, BatchingSink, SinkBackpressureError
 from quixstreams.sinks.base import SinkBatch
 from quixstreams.state import StateStoreManager
+from quixstreams.state.manager import SUPPORTED_STORES
 from quixstreams.state.base import PartitionTransaction
 from quixstreams.state.exceptions import StoreNotRegisteredError, StoreTransactionFailed
 from tests.utils import DummySink
@@ -71,6 +72,7 @@ def write(self, batch: SinkBatch):
         raise ValueError("Sink write failed")
 
 
+@pytest.mark.parametrize("store_type", SUPPORTED_STORES, indirect=True)
 class TestCheckpoint:
     def test_empty_true(self, checkpoint_factory):
         checkpoint = checkpoint_factory()

tests/test_quixstreams/test_state/fixtures.py

@@ -1,13 +1,23 @@
 import uuid
-from typing import Optional
-from unittest.mock import MagicMock
+from typing import Optional, Generator
+from unittest.mock import MagicMock, PropertyMock
 
 import pytest
+from quixstreams.state.base import StorePartition
 
 from quixstreams.kafka import Consumer
 from quixstreams.models import TopicManager
-from quixstreams.state.recovery import RecoveryPartition, RecoveryManager
-from quixstreams.state.base import StorePartition
+from quixstreams.state.recovery import (
+    RecoveryPartition,
+    RecoveryManager,
+    ChangelogProducerFactory,
+    ChangelogProducer,
+)
+from quixstreams.state.rocksdb import (
+    RocksDBStore,
+    RocksDBStorePartition,
+    RocksDBOptions,
+)
 
 
 @pytest.fixture()
@@ -45,3 +55,83 @@ def factory(
         return recovery_partition
 
     return factory
+
+
+@pytest.fixture()
+def store_type(request):
+    if hasattr(request, "param"):
+        return request.param
+    else:
+        return RocksDBStore
+
+
+def rocksdb_store_factory(tmp_path):
+    def factory(
+        topic: Optional[str] = None,
+        name: str = "default",
+        changelog_producer_factory: Optional[ChangelogProducerFactory] = None,
+    ) -> RocksDBStore:
+        topic = topic or str(uuid.uuid4())
+        return RocksDBStore(
+            topic=topic,
+            name=name,
+            base_dir=str(tmp_path),
+            changelog_producer_factory=changelog_producer_factory,
+        )
+
+    return factory
+
+
+@pytest.fixture()
+def store_factory(store_type, tmp_path):
+    if store_type == RocksDBStore:
+        return rocksdb_store_factory(tmp_path)
+    else:
+        raise ValueError(f"invalid store type {store_type}")
+
+
+@pytest.fixture()
+def store(store_factory):
+    store = store_factory()
+    yield store
+    store.close()
+
+
+def rocksdb_partition_factory(tmp_path, changelog_producer_mock):
+    def factory(
+        name: str = "db",
+        options: Optional[RocksDBOptions] = None,
+        changelog_producer: Optional[ChangelogProducer] = None,
+    ) -> RocksDBStorePartition:
+        path = (tmp_path / name).as_posix()
+        _options = options or RocksDBOptions(open_max_retries=0, open_retry_backoff=3.0)
+        return RocksDBStorePartition(
+            path,
+            changelog_producer=changelog_producer or changelog_producer_mock,
+            options=_options,
+        )
+
+    return factory
+
+
+@pytest.fixture()
+def store_partition_factory(store_type, tmp_path, changelog_producer_mock):
+    if store_type == RocksDBStore:
+        return rocksdb_partition_factory(tmp_path, changelog_producer_mock)
+    else:
+        raise ValueError(f"invalid store type {store_type}")
+
+
+@pytest.fixture()
+def store_partition(store_partition_factory) -> Generator[StorePartition, None, None]:
+    partition = store_partition_factory()
+    yield partition
+    partition.close()
+
+
+@pytest.fixture()
+def changelog_producer_mock():
+    producer = MagicMock(spec_set=ChangelogProducer)
+    type(producer).changelog_name = PropertyMock(return_value="test-changelog-topic")
+    type(producer).partition = PropertyMock(return_value=0)
+    return producer