Make PartitionTransaction.begin() abstract

Author: daniil-quix

quixstreams/state/base/partition.py

@@ -1,6 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Literal, Optional, Type, Union
+from typing import TYPE_CHECKING, Literal, Optional, Union
 
 from quixstreams.state.metadata import (
     Marker,
@@ -25,8 +25,6 @@ class StorePartition(ABC):
     the persistent storage).
     """
 
-    partition_transaction_class: Type[PartitionTransaction] = PartitionTransaction
-
     def __init__(
         self,
         dumps: DumpsFunc,
@@ -108,18 +106,14 @@ def recover_from_changelog_message(
         :param offset: changelog message offset
         """
 
+    @abstractmethod
     def begin(self) -> PartitionTransaction:
         """
         Start a new `PartitionTransaction`
 
         Using `PartitionTransaction` is a recommended way for accessing the data.
         """
-        return self.partition_transaction_class(
-            partition=self,
-            dumps=self._dumps,
-            loads=self._loads,
-            changelog_producer=self._changelog_producer,
-        )
+        ...
 
     def __enter__(self):
         return self

quixstreams/state/base/transaction.py

@@ -44,6 +44,7 @@
     "PartitionTransactionStatus",
     "PartitionTransaction",
     "PartitionTransactionCache",
+    "validate_transaction_status",
 )
 
 logger = logging.getLogger(__name__)

quixstreams/state/memory/partition.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Any, Dict, Literal, Optional, Union
 
+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
@@ -62,6 +63,14 @@ def closed(self) -> bool:
     def close(self) -> None:
         self._closed = True
 
+    def begin(self) -> PartitionTransaction:
+        return PartitionTransaction(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )
+
     @_validate_partition_state()
     def write(
         self,

quixstreams/state/rocksdb/partition.py

@@ -13,6 +13,7 @@
 from rocksdict import AccessType, ColumnFamily, Rdict, ReadOptions, WriteBatch
 
 from quixstreams.state.base import (
+    PartitionTransaction,
     PartitionTransactionCache,
     StorePartition,
 )
@@ -199,6 +200,14 @@ def iter_items(
                 if lower_bound <= key:
                     yield key, value
 
+    def begin(self) -> PartitionTransaction:
+        return PartitionTransaction(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )
+
     def exists(self, key: bytes, cf_name: str = "default") -> bool:
         """
         Check if a key is present in the DB.

quixstreams/state/rocksdb/timestamped.py

@@ -240,9 +240,17 @@ class TimestampedStorePartition(RocksDBStorePartition):
     `TimestampedPartitionTransaction` instances to handle atomic operations for that partition.
     """
 
-    partition_transaction_class = TimestampedPartitionTransaction
     additional_column_families = (MIN_ELIGIBLE_TIMESTAMPS_CF_NAME,)
 
+    def begin(self) -> TimestampedPartitionTransaction:
+        # Override the method to specify the correct return type
+        return TimestampedPartitionTransaction(
+            partition=self,
+            dumps=self._dumps,
+            loads=self._loads,
+            changelog_producer=self._changelog_producer,
+        )
+
 
 class TimestampedStore(RocksDBStore):
     """

quixstreams/state/rocksdb/windowed/partition.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterator, Type, cast
+from typing import Iterator, cast
 
 from ..partition import RocksDBStorePartition
 from .metadata import (
@@ -24,9 +24,6 @@ class WindowedRocksDBStorePartition(RocksDBStorePartition):
     stores the expiration index to delete expired windows.
     """
 
-    partition_transaction_class: Type[WindowedRocksDBPartitionTransaction] = (
-        WindowedRocksDBPartitionTransaction
-    )
     additional_column_families = (
         LATEST_DELETED_VALUE_CF_NAME,
         LATEST_EXPIRED_WINDOW_CF_NAME,
@@ -54,7 +51,7 @@ def begin(self) -> WindowedRocksDBPartitionTransaction:
 
         Using `WindowedRocksDBPartitionTransaction` is a recommended way for accessing the data.
         """
-        return self.partition_transaction_class(
+        return WindowedRocksDBPartitionTransaction(
             partition=self,
             dumps=self._dumps,
             loads=self._loads,