Handle various corruption scenarios

Author: gwaramadze

quixstreams/state/rocksdb/partition.py

@@ -1,5 +1,4 @@
 import logging
-import re
 import time
 from typing import (
     Dict,
@@ -35,12 +34,6 @@
 
 logger = logging.getLogger(__name__)
 
-ROCKSDB_CORRUPTION_REGEX = re.compile(
-    r"^Corruption: Corruption: IO error: No such file or "
-    r"directory: While open a file for random read: .*: "
-    r"No such file or directory  The file .* may be corrupted\.$"
-)
-
 
 class RocksDBStorePartition(StorePartition):
     """
@@ -366,7 +359,7 @@ def _open_rocksdict(self) -> Rdict:
         try:
             rdict = create_rdict()
         except Exception as exc:
-            if ROCKSDB_CORRUPTION_REGEX.fullmatch(str(exc)) is None:
+            if not str(exc).startswith("Corruption"):
                 raise
             elif not self._changelog_producer:
                 raise RocksDBCorruptedError(

tests/test_quixstreams/test_state/test_rocksdb/test_rocksdb_partition.py

@@ -16,19 +16,6 @@
 from quixstreams.state.rocksdb.windowed.serialization import append_integer
 
 
-@pytest.fixture
-def corrupted_db(tmp_path):
-    """
-    Create a corrupted RocksDB database in a temporary directory.
-    """
-    rdict = Rdict(path=tmp_path.as_posix())
-    # Write something and flush so that the .sst file is created
-    rdict[b"key"] = b"value"
-    rdict.flush()
-    # Delete the .sst file
-    next(tmp_path.glob("*.sst")).unlink()
-
-
 class TestRocksDBStorePartition:
     def test_open_db_locked_retries(self, store_partition_factory, executor):
         db1 = store_partition_factory("db")
@@ -83,8 +70,12 @@ def test_open_arbitrary_exception_fails(self, store_partition_factory):
         assert str(raised.value) == "some exception"
 
     def test_db_corrupted_fails_with_no_changelog(
-        self, store_partition_factory, corrupted_db
+        self, store_partition_factory, tmp_path
     ):
+        # Initialize and corrupt the database by messing with the MANIFEST
+        Rdict(path=tmp_path.as_posix())
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")
+
         with pytest.raises(
             RocksDBCorruptedError,
             match=(
@@ -95,8 +86,12 @@ def test_db_corrupted_fails_with_no_changelog(
             store_partition_factory(changelog_producer=None)
 
     def test_db_corrupted_fails_with_on_corrupted_recreate_false(
-        self, store_partition_factory, corrupted_db
+        self, store_partition_factory, tmp_path
     ):
+        # Initialize and corrupt the database by messing with the MANIFEST
+        Rdict(path=tmp_path.as_posix())
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")
+
         with pytest.raises(
             RocksDBCorruptedError,
             match=re.escape(
@@ -107,7 +102,19 @@ def test_db_corrupted_fails_with_on_corrupted_recreate_false(
         ):
             store_partition_factory()
 
-    def test_db_corrupted_recreated(self, store_partition_factory, corrupted_db):
+    def test_db_corrupted_manifest_file(self, store_partition_factory, tmp_path):
+        Rdict(path=tmp_path.as_posix())  # initialize db
+        next(tmp_path.glob("MANIFEST*")).write_bytes(b"")  # write random bytes
+
+        store_partition_factory(options=RocksDBOptions(on_corrupted_recreate=True))
+
+    def test_db_corrupted_sst_file(self, store_partition_factory, tmp_path):
+        rdict = Rdict(path=tmp_path.as_posix())  # initialize db
+        rdict[b"key"] = b"value"  # write something
+        rdict.flush()  # flush creates .sst file
+        rdict.close()  # required to release the lock
+        next(tmp_path.glob("*.sst")).unlink()  # delete the .sst file
+
         store_partition_factory(options=RocksDBOptions(on_corrupted_recreate=True))
 
     def test_create_and_get_column_family(self, store_partition: RocksDBStorePartition):