Check partitions status after committing in checkpoint

Author: ovv

quixstreams/checkpointing/checkpoint.py

@@ -2,7 +2,7 @@
 import time
 from typing import Dict, Tuple
 
-from confluent_kafka import TopicPartition
+from confluent_kafka import TopicPartition, KafkaException
 
 from quixstreams.kafka import Consumer
 from quixstreams.rowproducer import RowProducer
@@ -12,7 +12,11 @@
     DEFAULT_STATE_STORE_NAME,
 )
 from quixstreams.state.exceptions import StoreTransactionFailed
-from .exceptions import InvalidStoredOffset, CheckpointProducerTimeout
+from .exceptions import (
+    InvalidStoredOffset,
+    CheckpointProducerTimeout,
+    CheckpointConsumerCommitError,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -145,7 +149,14 @@ def commit(self):
             for (topic, partition), offset in self._tp_offsets.items()
         ]
         logger.debug("Checkpoint: commiting consumer")
-        self._consumer.commit(offsets=offsets, asynchronous=False)
+        try:
+            partitions = self._consumer.commit(offsets=offsets, asynchronous=False)
+        except KafkaException as e:
+            raise CheckpointConsumerCommitError(e.args[0]) from None
+
+        for partition in partitions:
+            if partition.error:
+                raise CheckpointConsumerCommitError(partition.error)
 
         # Step 4. Flush state store partitions to the disk together with changelog
         # offsets

quixstreams/checkpointing/exceptions.py

@@ -1,7 +1,11 @@
 from quixstreams.exceptions import QuixException
+from quixstreams.kafka.exceptions import KafkaConsumerException
 
 
 class InvalidStoredOffset(QuixException): ...
 
 
 class CheckpointProducerTimeout(QuixException): ...
+
+
+class CheckpointConsumerCommitError(KafkaConsumerException): ...

tests/test_quixstreams/test_checkpointing.py

@@ -3,10 +3,13 @@
 from unittest.mock import patch, MagicMock
 
 import pytest
-from confluent_kafka import TopicPartition
+from confluent_kafka import TopicPartition, KafkaException, KafkaError
 
 from quixstreams.checkpointing import Checkpoint, InvalidStoredOffset
-from quixstreams.checkpointing.exceptions import CheckpointProducerTimeout
+from quixstreams.checkpointing.exceptions import (
+    CheckpointProducerTimeout,
+    CheckpointConsumerCommitError,
+)
 from quixstreams.kafka import Consumer
 from quixstreams.rowproducer import RowProducer
 from quixstreams.state import StateStoreManager
@@ -354,3 +357,50 @@ def test_incomplete_flush(
             str(err.value)
             == "'1' messages failed to be produced before the producer flush timeout"
         )
+
+    def test_failed_commit(
+        self, checkpoint_factory, state_manager_factory, rowproducer_mock
+    ):
+        consumer_mock = MagicMock(spec_set=Consumer)
+        consumer_mock.commit.side_effect = KafkaException(KafkaError(1, "test error"))
+
+        state_manager = state_manager_factory(producer=rowproducer_mock)
+        checkpoint = checkpoint_factory(
+            consumer_=consumer_mock,
+            state_manager_=state_manager,
+            producer_=rowproducer_mock,
+        )
+        checkpoint.store_offset("topic", 0, 0)
+
+        with pytest.raises(CheckpointConsumerCommitError) as err:
+            checkpoint.commit()
+
+        assert (
+            str(err.value)
+            == '<CheckpointConsumerCommitError code="1" description="test error">'
+        )
+
+    def test_failed_commit_partition(
+        self, checkpoint_factory, state_manager_factory, rowproducer_mock
+    ):
+        consumer_mock = MagicMock(spec_set=Consumer)
+
+        topic_partition = MagicMock(spec=TopicPartition)
+        topic_partition.error = KafkaError(1, "test error")
+
+        consumer_mock.commit.return_value = [topic_partition]
+        state_manager = state_manager_factory(producer=rowproducer_mock)
+        checkpoint = checkpoint_factory(
+            consumer_=consumer_mock,
+            state_manager_=state_manager,
+            producer_=rowproducer_mock,
+        )
+        checkpoint.store_offset("topic", 0, 0)
+
+        with pytest.raises(CheckpointConsumerCommitError) as err:
+            checkpoint.commit()
+
+        assert (
+            str(err.value)
+            == '<CheckpointConsumerCommitError code="1" description="test error">'
+        )