Add a method to clear the application local state (#239)

Add Application.clear_state() method to clear the local state directory before the process starts.

Author: harisbotic

src/StreamingDataFrames/docs/stateful-processing.md

@@ -73,6 +73,19 @@ Currently, only functions passed to `StreamingDataFrame.apply()`, `StreamingData
 
 <br>
 
+## Clearing the State
+
+To clear all the state data, use the `app.clear_state()` command. This will delete all data stored in the state stores, allowing you to start from a clean slate:
+
+```python
+app.clear_state()
+```
+
+Note that clearing the app state using `app.clear_state()` is only possible when the `app.run()` is not running. Meaning that the state can be cleared either before calling `app.run()` or after.
+This ensures that state clearing does not interfere with the ongoing stateful processing.
+
+<br>
+
 ## Changing the State FilePath
 
 By default, an `Application` keeps the state in `state` directory relative to the current working directory.

src/StreamingDataFrames/quixstreams/app.py

@@ -346,6 +346,12 @@ def stop(self):
         """
         self._running = False
 
+    def clear_state(self):
+        """
+        Clear the state of the application.
+        """
+        self._state_manager.clear_stores()
+
     def _quix_runtime_init(self):
         """
         Do a runtime setup only applicable to an Application.Quix instance

src/StreamingDataFrames/quixstreams/state/exceptions.py

@@ -5,6 +5,10 @@ class PartitionNotAssignedError(QuixException):
     ...
 
 
+class PartitionStoreIsUsed(QuixException):
+    ...
+
+
 class StoreNotRegisteredError(QuixException):
     ...
 

src/StreamingDataFrames/quixstreams/state/manager.py

@@ -1,12 +1,14 @@
 import contextlib
 import logging
+import shutil
 from pathlib import Path
 from typing import List, Dict, Optional, Iterator
 
 from quixstreams.types import TopicPartition
 from .exceptions import (
     StoreNotRegisteredError,
     InvalidStoreTransactionStateError,
+    PartitionStoreIsUsed,
 )
 from .rocksdb import RocksDBStore, RocksDBOptionsType
 from .types import (
@@ -104,6 +106,21 @@ def register_store(
                 options=self._rocksdb_options,
             )
 
+    def clear_stores(self):
+        """
+        Delete all state stores managed by StateStoreManager.
+        """
+        if any(
+            store.partitions
+            for topic_stores in self._stores.values()
+            for store in topic_stores.values()
+        ):
+            raise PartitionStoreIsUsed(
+                "Cannot clear stores with active partitions assigned"
+            )
+
+        shutil.rmtree(self._state_dir)
+
     def on_partition_assign(self, tp: TopicPartition) -> List[StorePartition]:
         """
         Assign store partitions for each registered store for the given `TopicPartition`

src/StreamingDataFrames/tests/test_quixstreams/test_app.py

@@ -717,3 +717,58 @@ def test_on_assign_topic_offset_behind_warning(
 
         assert mock.called
         assert "is behind the stored offset" in mock.call_args[0][0]
+
+    def test_clear_state(
+        self,
+        app_factory,
+        producer,
+        topic_factory,
+        executor,
+        state_manager_factory,
+        tmp_path,
+    ):
+        """
+        Test that clear_state() removes all data from the state stores
+        """
+
+        consumer_group = str(uuid.uuid4())
+        state_dir = (tmp_path / "state").absolute()
+        app = app_factory(
+            consumer_group=consumer_group,
+            auto_offset_reset="earliest",
+            state_dir=state_dir,
+        )
+
+        topic_in_name, _ = topic_factory()
+        tx_prefix = b"key"
+
+        state_manager = state_manager_factory(
+            group_id=consumer_group, state_dir=state_dir
+        )
+
+        # Add data to the state store
+        with state_manager:
+            state_manager.register_store(topic_in_name, "default")
+            state_manager.on_partition_assign(
+                TopicPartitionStub(topic=topic_in_name, partition=0)
+            )
+            store = state_manager.get_store(topic=topic_in_name, store_name="default")
+            with store.start_partition_transaction(partition=0) as tx:
+                # All keys in state must be prefixed with the message key
+                with tx.with_prefix(tx_prefix):
+                    tx.set("my_state", True)
+
+        # Clear the state
+        app.clear_state()
+
+        # Check that the date is cleared from the state store
+        with state_manager:
+            state_manager.register_store(topic_in_name, "default")
+            state_manager.on_partition_assign(
+                TopicPartitionStub(topic=topic_in_name, partition=0)
+            )
+            store = state_manager.get_store(topic=topic_in_name, store_name="default")
+            with store.start_partition_transaction(partition=0) as tx:
+                # All keys in state must be prefixed with the message key
+                with tx.with_prefix(tx_prefix):
+                    assert tx.get("my_state") is None

src/StreamingDataFrames/tests/test_quixstreams/test_state/test_manager.py

@@ -1,3 +1,4 @@
+import os
 import contextlib
 import uuid
 from unittest.mock import patch
@@ -9,6 +10,7 @@
 from quixstreams.state.exceptions import (
     StoreNotRegisteredError,
     InvalidStoreTransactionStateError,
+    PartitionStoreIsUsed,
 )
 
 
@@ -108,6 +110,60 @@ def test_get_store_not_registered(self, state_manager):
         with pytest.raises(StoreNotRegisteredError):
             state_manager.get_store("topic", "store")
 
+    def test_clear_stores_when_empty(self, state_manager):
+        state_manager.clear_stores()
+        assert not state_manager.stores
+
+    def test_clear_stores(self, state_manager):
+        # Register stores
+        state_manager.register_store("topic1", store_name="store1")
+        state_manager.register_store("topic1", store_name="extra_store")
+        state_manager.register_store("topic2", store_name="store1")
+
+        # Define partitions
+        partitions = [
+            TopicPartitionStub("topic1", 0),
+            TopicPartitionStub("topic1", 1),
+            TopicPartitionStub("topic2", 0),
+        ]
+
+        # Assign partitions
+        for tp in partitions:
+            state_manager.on_partition_assign(tp)
+
+        # Collect paths of stores to be deleted
+        stores_to_delete = [
+            store_partition.path
+            for topic_stores in state_manager.stores.values()
+            for store in topic_stores.values()
+            for store_partition in store.partitions.values()
+        ]
+
+        # Revoke partitions
+        for tp in partitions:
+            state_manager.on_partition_revoke(tp)
+
+        # Act - Delete stores
+        state_manager.clear_stores()
+
+        # Assert store paths are deleted
+        for path in stores_to_delete:
+            assert not os.path.exists(path), f"RocksDB store at {path} was not deleted"
+
+    def test_clear_stores_fails(self, state_manager):
+        # Register stores
+        state_manager.register_store("topic1", store_name="store1")
+
+        # Define the partition
+        partition = TopicPartitionStub("topic1", 0)
+
+        # Assign the partition
+        state_manager.on_partition_assign(partition)
+
+        # Act - Delete stores
+        with pytest.raises(PartitionStoreIsUsed):
+            state_manager.clear_stores()
+
     def test_store_transaction_success(self, state_manager):
         state_manager.register_store("topic", "store")
         tp = TopicPartitionStub("topic", 0)