Call prepare on PartitionTransaction context manager exit

Author: gwaramadze

quixstreams/sources/base/source.py

@@ -528,7 +528,7 @@ def flush(self, timeout: Optional[float] = None) -> None:
         :raises CheckpointProducerTimeout: if any message fails to produce before the timeout
         """
         if self._store_transaction:
-            self._store_transaction.prepare(None)
+            self._store_transaction.prepare()
             self._store_transaction.flush()
             self._store_transaction = None
             self._store_state = None

quixstreams/state/base/transaction.py

@@ -471,7 +471,7 @@ def exists(self, key: K, prefix: bytes, cf_name: str = "default") -> bool:
             return self._partition.exists(key_serialized, cf_name=cf_name)
 
     @validate_transaction_status(PartitionTransactionStatus.STARTED)
-    def prepare(self, processed_offsets: Optional[dict[str, int]]):
+    def prepare(self, processed_offsets: Optional[dict[str, int]] = None) -> None:
         """
         Produce changelog messages to the changelog topic for all changes accumulated
         in this transaction and prepare transaction to flush its state to the state
@@ -581,4 +581,5 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if exc_val is None and not self.failed:
+            self.prepare()
             self.flush()