"range" assignment updates (#712)

* Don't commit the checkpoint on assign
* Enforce "range" assignor for internal consumers

Author: daniil-quix

quixstreams/app.py

@@ -58,6 +58,9 @@
 # Enforce idempotent producing for the internal RowProducer
 _default_producer_extra_config = {"enable.idempotence": True}
 
+# Force assignment strategy to be "range" for co-partitioning in internal Consumers
+consumer_extra_config_overrides = {"partition.assignment.strategy": "range"}
+
 _default_max_poll_interval_ms = 300000
 
 
@@ -301,7 +304,10 @@ def __init__(
         self._on_message_processed = on_message_processed
         self._on_processing_error = on_processing_error or default_on_processing_error
 
-        self._consumer = self._get_rowconsumer(on_error=on_consumer_error)
+        self._consumer = self._get_rowconsumer(
+            on_error=on_consumer_error,
+            extra_config_overrides=consumer_extra_config_overrides,
+        )
         self._producer = self._get_rowproducer(on_error=on_producer_error)
         self._running = False
         self._failed = False
@@ -593,20 +599,27 @@ def get_producer(self) -> Producer:
         )
 
     def _get_rowconsumer(
-        self, on_error: Optional[ConsumerErrorCallback] = None
+        self,
+        on_error: Optional[ConsumerErrorCallback] = None,
+        extra_config_overrides: Optional[dict] = None,
     ) -> RowConsumer:
         """
         Create a RowConsumer using the application config
 
         Used to create the application consumer as well as the sources consumers
         """
-
+        extra_config_overrides = extra_config_overrides or {}
+        # Override the existing "extra_config" with new values
+        extra_config = {
+            **self._config.consumer_extra_config,
+            **extra_config_overrides,
+        }
         return RowConsumer(
             broker_address=self._config.broker_address,
             consumer_group=self._config.consumer_group,
             auto_offset_reset=self._config.auto_offset_reset,
             auto_commit_enable=False,  # Disable auto commit and manage commits manually
-            extra_config=self._config.consumer_extra_config,
+            extra_config=extra_config,
             on_error=on_error,
         )
 
@@ -692,7 +705,9 @@ def add_source(self, source: BaseSource, topic: Optional[Topic] = None) -> Topic
             source,
             topic,
             self._get_rowproducer(transactional=False),
-            self._get_rowconsumer(),
+            self._get_rowconsumer(
+                extra_config_overrides=consumer_extra_config_overrides
+            ),
             self._get_topic_manager(),
         )
         return topic
@@ -897,10 +912,6 @@ def _on_assign(self, _, topic_partitions: List[TopicPartition]):
         # get the source data.
         self._source_manager.start_sources()
 
-        # First commit everything processed so far because assignment can take a while
-        # and fail
-        self._processing_context.commit_checkpoint(force=True)
-
         # Assign partitions manually to pause the changelog topics
         self._consumer.assign(topic_partitions)
         # Pause changelog topic+partitions immediately after assignment

quixstreams/kafka/consumer.py

@@ -126,7 +126,6 @@ def __init__(
         self._consumer_config = {
             # previous Quix Streams defaults
             "enable.auto.offset.store": False,
-            # Force assignment strategy to "range" for co-partitioning
             "partition.assignment.strategy": "range",
             **(extra_config or {}),
             **broker_address.as_librdkafka_dict(),

tests/test_quixstreams/test_app.py

@@ -839,6 +839,7 @@ def test_init_with_quix_sdk_token_arg(self):
         expected_consumer_extra_config = {
             **extra_config,
             **quix_extras,
+            "partition.assignment.strategy": "range",
         }
 
         def get_cfg_builder(quix_sdk_token):
@@ -903,6 +904,7 @@ def test_init_with_quix_sdk_token_env(self, monkeypatch):
         expected_consumer_extra_config = {
             **extra_config,
             **quix_extras,
+            "partition.assignment.strategy": "range",
         }
 
         def get_cfg_builder(quix_sdk_token):
@@ -965,6 +967,7 @@ def test_init_with_quix_config_builder(self):
         expected_consumer_extra_config = {
             **extra_config,
             **quix_extras,
+            "partition.assignment.strategy": "range",
         }
 
         def get_cfg_builder(quix_sdk_token):