external kafka source

Author: ovv

quixstreams/app.py

@@ -667,7 +667,9 @@ def get_producer(self) -> Producer:
             extra_config=self._producer_extra_config,
         )
 
-    def get_consumer(self, auto_commit_enable: bool = True) -> Consumer:
+    def get_consumer(
+        self, auto_commit_enable: bool = True, consumer_group_suffix: str = ""
+    ) -> Consumer:
         """
         Create and return a pre-configured Consumer instance.
         The Consumer is initialized with params passed to Application.
@@ -708,7 +710,7 @@ def get_consumer(self, auto_commit_enable: bool = True) -> Consumer:
 
         return Consumer(
             broker_address=self._broker_address,
-            consumer_group=self._consumer_group,
+            consumer_group=f"{self._consumer_group}{consumer_group_suffix}",
             auto_offset_reset=self._auto_offset_reset,
             auto_commit_enable=auto_commit_enable,
             extra_config=self._consumer_extra_config,

quixstreams/sources/kafka.py

@@ -0,0 +1,240 @@
+import logging
+import threading
+
+from typing import Union, Optional, Dict, TYPE_CHECKING
+
+from confluent_kafka import TopicPartition, KafkaException
+
+from quixstreams.app import Application
+from quixstreams.kafka import Consumer, ConnectionConfig, AutoOffsetReset
+from quixstreams.kafka.exceptions import KafkaConsumerException
+from quixstreams.platforms.quix import QuixKafkaConfigsBuilder
+from quixstreams.checkpointing.exceptions import (
+    CheckpointProducerTimeout,
+    CheckpointConsumerCommitError,
+)
+from quixstreams.models.topics import TopicConfig, TopicAdmin
+from quixstreams.platforms.quix.api import QuixPortalApiService
+
+from .base import CheckpointingSource, SourceTopic
+
+if TYPE_CHECKING:
+    from quixstreams.app import Application
+
+logger = logging.getLogger(__name__)
+
+
+class ExternalKafkaSource(CheckpointingSource):
+
+    def __init__(
+        self,
+        name: str,
+        app: "Application",
+        topic: str,
+        broker_address: Union[str, ConnectionConfig],
+        consumer_group: Optional[str] = None,
+        auto_offset_reset: AutoOffsetReset = "latest",
+        consumer_extra_config: Optional[dict] = None,
+        shutdown_timeout: int = 10,
+    ) -> None:
+        super().__init__(name, shutdown_timeout)
+
+        if consumer_extra_config is None:
+            consumer_extra_config = {}
+
+        if consumer_group is None:
+            consumer_group = name
+
+        self._source_cluster_consumer = Consumer(
+            broker_address=broker_address,
+            consumer_group=consumer_group,
+            auto_offset_reset=auto_offset_reset,
+            extra_config=consumer_extra_config,
+            auto_commit_enable=False,
+        )
+
+        self._source_cluster_admin = TopicAdmin(
+            broker_address=broker_address,
+            extra_config=consumer_extra_config,
+        )
+
+        self._target_cluster_consumer = app.get_consumer(
+            auto_commit_enable=False, consumer_group_suffix="-offsets"
+        )
+
+        self._running = True
+        self._topic = topic
+        self._tp_offsets = {}
+        self._tp_offsets_lock = threading.Lock()
+
+        self._error: Optional[Exception] = None
+
+    def run(self) -> None:
+        self._source_cluster_consumer.subscribe(
+            topics=[self._topic],
+            on_assign=self._on_assign,
+            on_lost=self._on_lost,
+            on_revoke=self._on_revoke,
+        )
+
+        super().run()
+        while self._running:
+            if self._error:
+                raise self._error
+
+            msg = self._source_cluster_consumer.poll(timeout=1)
+            if msg is None:
+                continue
+
+            if err := msg.error():
+                raise KafkaConsumerException(error=err)
+
+            topic_name, partition, offset = msg.topic(), msg.partition(), msg.offset()
+            self.produce(
+                value=msg.value(),
+                key=msg.key(),
+                headers=msg.headers(),
+                timestamp=msg.timestamp()[1],
+                partition=partition,
+            )
+            with self._tp_offsets_lock:
+                self._tp_offsets[(topic_name, partition)] = offset
+
+    def _target_cluster_offsets(self, partitions) -> Dict[int, int]:
+        partitions = [
+            TopicPartition(
+                topic=self._producer_topic.name, partition=partition.partition
+            )
+            for partition in partitions
+        ]
+        partitions_commited = self._target_cluster_consumer.committed(
+            partitions, timeout=5
+        )
+        return {
+            partition.partition: partition.offset for partition in partitions_commited
+        }
+
+    def _on_assign(self, _, source_partitions):
+        try:
+            target_cluster_offset = self._target_cluster_offsets(source_partitions)
+            for partition in source_partitions:
+                partition.offset = target_cluster_offset.get(partition.partition, None)
+
+            self._source_cluster_consumer.incremental_assign(source_partitions)
+        except Exception as exc:
+            logger.exception("Error in source %s assigning partitions", self)
+            self._error = exc
+
+    def _on_revoke(self, _, partitions):
+        if not self._error:
+            self.checkpoint()
+
+    def _on_lost(self, _, partitions):
+        pass
+
+    def checkpoint(self):
+        if not self._tp_offsets:
+            return
+
+        super().checkpoint()
+        with self._tp_offsets_lock:
+            offsets = [
+                TopicPartition(
+                    topic=self._producer_topic.name,
+                    partition=partition,
+                    offset=offset + 1,
+                )
+                for (topic, partition), offset in self._tp_offsets.items()
+            ]
+            self._tp_offsets = {}
+
+        unproduced_msg_count = self._producer.flush(self.shutdown_timeout / 2)
+        if unproduced_msg_count > 0:
+            raise CheckpointProducerTimeout(
+                f"'{unproduced_msg_count}' messages failed to be produced before the producer flush timeout"
+            )
+
+        try:
+            partitions = self._target_cluster_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)
+
+    def stop(self):
+        super().stop()
+        self._running = False
+
+    def cleanup(self, failed):
+        super().cleanup(failed)
+        if not failed:
+            self.checkpoint()
+
+        self._source_cluster_consumer.close()
+        self._target_cluster_consumer.close()
+
+    def default_topic(self) -> SourceTopic:
+        config = self._source_cluster_admin.inspect_topics(
+            topic_names=[self._topic], timeout=5
+        ).get(self._topic)
+        if config is None:
+            config = TopicConfig(num_partitions=1, replication_factor=1)
+
+        return SourceTopic(
+            name=self.name,
+            value_serializer="bytes",
+            value_deserializer="json",
+            config=config,
+        )
+
+
+class QuixEnvironmentSource(ExternalKafkaSource):
+    def __init__(
+        self,
+        name: str,
+        app: "Application",
+        topic: str,
+        quix_sdk_token: str,
+        quix_workspace_id: str,
+        consumer_group: Optional[str] = None,
+        quix_portal_api: Optional[str] = None,
+        auto_offset_reset: AutoOffsetReset = "latest",
+        consumer_extra_config: Optional[dict] = None,
+        shutdown_timeout: int = 10,
+    ) -> None:
+
+        if consumer_extra_config is None:
+            consumer_extra_config = {}
+
+        if consumer_group is None:
+            consumer_group = name
+
+        quix_config = QuixKafkaConfigsBuilder(
+            quix_portal_api_service=QuixPortalApiService(
+                default_workspace_id=quix_workspace_id,
+                auth_token=quix_sdk_token,
+                portal_api=quix_portal_api,
+            )
+        )
+
+        self._short_topic = topic
+        self._quix_workspace_id = quix_workspace_id
+
+        consumer_extra_config.update(quix_config.librdkafka_extra_config)
+        super().__init__(
+            name=name,
+            app=app,
+            topic=quix_config.prepend_workspace_id(topic),
+            consumer_group=quix_config.prepend_workspace_id(consumer_group),
+            auto_offset_reset=auto_offset_reset,
+            consumer_extra_config=consumer_extra_config,
+            shutdown_timeout=shutdown_timeout,
+            broker_address=quix_config.librdkafka_connection_config,
+        )
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}({self._short_topic}@{self._quix_workspace_id})>"