sentry-arroyo 2.30.0__tar.gz → 2.32.0__tar.gz

{sentry_arroyo-2.30.0/sentry_arroyo.egg-info → sentry_arroyo-2.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.30.0
+Version: 2.32.0
 Summary: Arroyo is a Python library for working with streaming data.
 Home-page: https://github.com/getsentry/arroyo
 Author: Sentry

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/backends/abstract.py

@@ -153,10 +153,13 @@ class Consumer(Generic[TStrategyPayload], ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def commit_offsets(self) -> Mapping[Partition, int]:
+    def commit_offsets(self) -> Optional[Mapping[Partition, int]]:
         """
         Commit staged offsets. The return value of this method is a mapping
         of streams with their committed offsets as values.
+
+        When auto-commit is enabled (in Kafka consumers), returns None since
+        the broker handles commits automatically.
         """
         raise NotImplementedError
 

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/backends/kafka/__init__.py

@@ -3,12 +3,13 @@ from .configuration import (
     build_kafka_consumer_configuration,
     build_kafka_producer_configuration,
 )
-from .consumer import KafkaConsumer, KafkaPayload, KafkaProducer
+from .consumer import ConfluentProducer, KafkaConsumer, KafkaPayload, KafkaProducer
 
 __all__ = [
     "build_kafka_configuration",
     "build_kafka_consumer_configuration",
     "build_kafka_producer_configuration",
+    "ConfluentProducer",
     "KafkaConsumer",
     "KafkaPayload",
     "KafkaProducer",

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/backends/kafka/configuration.py

@@ -237,6 +237,7 @@ def build_kafka_consumer_configuration(
     bootstrap_servers: Optional[Sequence[str]] = None,
     override_params: Optional[Mapping[str, Any]] = None,
     strict_offset_reset: Optional[bool] = None,
+    enable_auto_commit: bool = False,
 ) -> KafkaBrokerConfig:
 
     if auto_offset_reset is None:
@@ -252,20 +253,23 @@ def build_kafka_consumer_configuration(
         default_config, bootstrap_servers, override_params
     )
 
-    broker_config.update(
-        {
-            "enable.auto.commit": False,
-            "enable.auto.offset.store": False,
-            "group.id": group_id,
-            "auto.offset.reset": auto_offset_reset,
-            # this is an arroyo specific flag that only affects the consumer.
-            "arroyo.strict.offset.reset": strict_offset_reset,
-            # overridden to reduce memory usage when there's a large backlog
-            "queued.max.messages.kbytes": queued_max_messages_kbytes,
-            "queued.min.messages": queued_min_messages,
-            "enable.partition.eof": False,
-            "statistics.interval.ms": STATS_COLLECTION_FREQ_MS,
-            "stats_cb": stats_callback,
-        }
-    )
+    # Default configuration with manual commit management
+    config_update = {
+        "enable.auto.commit": False,
+        "enable.auto.offset.store": False,
+        "group.id": group_id,
+        "auto.offset.reset": auto_offset_reset,
+        # this is an arroyo specific flag that only affects the consumer.
+        "arroyo.strict.offset.reset": strict_offset_reset,
+        # this is an arroyo specific flag to enable auto-commit mode
+        "arroyo.enable.auto.commit": enable_auto_commit,
+        # overridden to reduce memory usage when there's a large backlog
+        "queued.max.messages.kbytes": queued_max_messages_kbytes,
+        "queued.min.messages": queued_min_messages,
+        "enable.partition.eof": False,
+        "statistics.interval.ms": STATS_COLLECTION_FREQ_MS,
+        "stats_cb": stats_callback,
+    }
+
+    broker_config.update(config_update)
     return broker_config

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/backends/kafka/consumer.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import logging
+import time
+from collections import defaultdict
 from concurrent.futures import Future
 from datetime import datetime
 from enum import Enum
@@ -21,6 +23,7 @@ from typing import (
     Tuple,
     Type,
     Union,
+    cast,
 )
 
 from confluent_kafka import (
@@ -33,7 +36,7 @@ from confluent_kafka import (
 from confluent_kafka import Consumer as ConfluentConsumer
 from confluent_kafka import KafkaError, KafkaException
 from confluent_kafka import Message as ConfluentMessage
-from confluent_kafka import Producer as ConfluentProducer
+from confluent_kafka import Producer as ConfluentKafkaProducer
 from confluent_kafka import TopicPartition as ConfluentTopicPartition
 
 from arroyo.backends.abstract import (
@@ -50,6 +53,7 @@ from arroyo.errors import (
 )
 from arroyo.types import BrokerValue, Partition, Topic
 from arroyo.utils.concurrent import execute
+from arroyo.utils.metrics import get_metrics
 from arroyo.utils.retries import BasicRetryPolicy
 
 logger = logging.getLogger(__name__)
@@ -182,6 +186,13 @@ class KafkaConsumer(Consumer[KafkaPayload]):
         if self.__strict_offset_reset is None:
             self.__strict_offset_reset = True
 
+        # Feature flag to enable rdkafka auto-commit with store_offsets
+        # When enabled, offsets are stored via store_offsets() and rdkafka
+        # automatically commits them periodically
+        self.__use_auto_commit = as_kafka_configuration_bool(
+            configuration.pop("arroyo.enable.auto.commit", False)
+        )
+
         if auto_offset_reset in {"smallest", "earliest", "beginning"}:
             self.__resolve_partition_starting_offset = (
                 self.__resolve_partition_offset_earliest
@@ -197,21 +208,32 @@ class KafkaConsumer(Consumer[KafkaPayload]):
         else:
             raise ValueError("invalid value for 'auto.offset.reset' configuration")
 
-        if (
-            as_kafka_configuration_bool(configuration.get("enable.auto.commit", "true"))
-            is not False
-        ):
-            raise ValueError("invalid value for 'enable.auto.commit' configuration")
+        # When auto-commit is disabled (default), we require explicit configuration
+        # When auto-commit is enabled, we allow rdkafka to handle commits
+        if not self.__use_auto_commit:
+            if (
+                as_kafka_configuration_bool(
+                    configuration.get("enable.auto.commit", "true")
+                )
+                is not False
+            ):
+                raise ValueError("invalid value for 'enable.auto.commit' configuration")
 
-        if (
-            as_kafka_configuration_bool(
-                configuration.get("enable.auto.offset.store", "true")
-            )
-            is not False
-        ):
-            raise ValueError(
-                "invalid value for 'enable.auto.offset.store' configuration"
-            )
+            if (
+                as_kafka_configuration_bool(
+                    configuration.get("enable.auto.offset.store", "true")
+                )
+                is not False
+            ):
+                raise ValueError(
+                    "invalid value for 'enable.auto.offset.store' configuration"
+                )
+        else:
+            # In auto-commit mode, enable auto.commit and keep auto.offset.store disabled
+            # We'll use store_offsets() manually to control which offsets get committed
+            configuration["enable.auto.commit"] = True
+            configuration["enable.auto.offset.store"] = False
+            configuration["on_commit"] = self.__on_commit_callback
 
         # NOTE: Offsets are explicitly managed as part of the assignment
         # callback, so preemptively resetting offsets is not enabled when
@@ -231,6 +253,19 @@ class KafkaConsumer(Consumer[KafkaPayload]):
 
         self.__state = KafkaConsumerState.CONSUMING
 
+    def __on_commit_callback(
+        self,
+        error: Optional[KafkaException],
+        partitions: Sequence[ConfluentTopicPartition],
+    ) -> None:
+        if error:
+            partition_info = [f"{p.topic}:{p.partition}" for p in partitions]
+            logger.warning(
+                "Commit failed: %s. Partitions: %s",
+                error,
+                partition_info,
+            )
+
     def __resolve_partition_offset_earliest(
         self, partition: ConfluentTopicPartition
     ) -> ConfluentTopicPartition:
@@ -568,7 +603,21 @@ class KafkaConsumer(Consumer[KafkaPayload]):
         # TODO: Maybe log a warning if these offsets exceed the current
         # offsets, since that's probably a side effect of an incorrect usage
         # pattern?
-        self.__staged_offsets.update(offsets)
+        if self.__use_auto_commit:
+            # When auto-commit is enabled, use store_offsets to stage offsets
+            # for rdkafka to auto-commit
+            if offsets:
+                self.__consumer.store_offsets(
+                    offsets=[
+                        ConfluentTopicPartition(
+                            partition.topic.name, partition.index, offset
+                        )
+                        for partition, offset in offsets.items()
+                    ]
+                )
+        else:
+            # Default behavior: manually track staged offsets
+            self.__staged_offsets.update(offsets)
 
     def __commit(self) -> Mapping[Partition, int]:
         if self.__state in {KafkaConsumerState.CLOSED, KafkaConsumerState.ERROR}:
@@ -616,15 +665,24 @@ class KafkaConsumer(Consumer[KafkaPayload]):
 
         return offsets
 
-    def commit_offsets(self) -> Mapping[Partition, int]:
+    def commit_offsets(self) -> Optional[Mapping[Partition, int]]:
         """
         Commit staged offsets for all partitions that this consumer is
         assigned to. The return value of this method is a mapping of
         partitions with their committed offsets as values.
 
+        When auto-commit is enabled, returns None since rdkafka handles
+        commits automatically and we don't track which offsets were committed.
+
         Raises an ``InvalidState`` if called on a closed consumer.
         """
-        return self.__commit_retry_policy.call(self.__commit)
+        if self.__use_auto_commit:
+            # When auto-commit is enabled, rdkafka commits automatically
+            # We don't track what was committed, so return None
+            # The offsets have already been staged via store_offsets()
+            return None
+        else:
+            return self.__commit_retry_policy.call(self.__commit)
 
     def close(self, timeout: Optional[float] = None) -> None:
         """
@@ -657,7 +715,7 @@ class KafkaProducer(Producer[KafkaPayload]):
         self, configuration: Mapping[str, Any], use_simple_futures: bool = False
     ) -> None:
         self.__configuration = configuration
-        self.__producer = ConfluentProducer(configuration)
+        self.__producer = ConfluentKafkaProducer(configuration)
         self.__shutdown_requested = Event()
 
         # The worker must execute in a separate thread to ensure that callbacks
@@ -742,3 +800,79 @@ class KafkaProducer(Producer[KafkaPayload]):
     def close(self) -> Future[None]:
         self.__shutdown_requested.set()
         return self.__result
+
+
+# Type alias for the delivery callback function
+DeliveryCallback = Callable[[Optional[KafkaError], ConfluentMessage], None]
+
+# Interval between metric flushes (in seconds)
+METRICS_FREQUENCY_SEC = 1.0
+
+
+class ConfluentProducer(ConfluentKafkaProducer):  # type: ignore[misc]
+    """
+    A thin wrapper for confluent_kafka.Producer that adds metrics reporting.
+    """
+
+    def __init__(self, configuration: Mapping[str, Any]) -> None:
+        super().__init__(configuration)
+        self.producer_name = configuration.get("client.id") or None
+        self.__metrics = get_metrics()
+        self.__produce_counters: MutableMapping[str, int] = defaultdict(int)
+        self.__reset_metrics()
+
+    def __metrics_delivery_callback(
+        self,
+        error: Optional[KafkaError],
+        _message: ConfluentMessage,
+    ) -> None:
+        if error is not None:
+            status = "error"
+        else:
+            status = "success"
+        self.__produce_counters[status] += 1
+        self.__throttled_record()
+
+    def __delivery_callback(
+        self,
+        user_callback: Optional[DeliveryCallback],
+    ) -> DeliveryCallback:
+        def wrapped(error: Optional[KafkaError], message: ConfluentMessage) -> None:
+            self.__metrics_delivery_callback(error, message)
+            if user_callback is not None:
+                user_callback(error, message)
+
+        return wrapped
+
+    def produce(self, *args: Any, **kwargs: Any) -> None:
+        # callback and on_delivery are aliases, callback takes precedence over on_delivery
+        callback = kwargs.pop("callback", None)
+        on_delivery = kwargs.pop("on_delivery", None)
+        user_callback = callback or on_delivery
+        wrapped_callback = self.__delivery_callback(user_callback)
+        super().produce(*args, on_delivery=wrapped_callback, **kwargs)
+
+    def __flush_metrics(self) -> None:
+        for status, count in self.__produce_counters.items():
+            tags = {"status": status}
+            if self.producer_name:
+                tags["producer_name"] = self.producer_name
+            self.__metrics.increment(
+                name="arroyo.producer.produce_status",
+                value=count,
+                tags=tags,
+            )
+        self.__reset_metrics()
+
+    def flush(self, timeout: float = -1) -> int:
+        # Kafka producer flush should flush metrics too
+        self.__flush_metrics()
+        return cast(int, super().flush(timeout))
+
+    def __reset_metrics(self) -> None:
+        self.__produce_counters.clear()
+        self.__last_record_time = time.time()
+
+    def __throttled_record(self) -> None:
+        if time.time() - self.__last_record_time > METRICS_FREQUENCY_SEC:
+            self.__flush_metrics()

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/backends/local/backend.py

@@ -38,9 +38,9 @@ class LocalBroker(Generic[TStrategyPayload]):
         self.__message_storage = message_storage
         self.__clock = clock
 
-        self.__offsets: MutableMapping[str, MutableMapping[Partition, int]] = (
-            defaultdict(dict)
-        )
+        self.__offsets: MutableMapping[
+            str, MutableMapping[Partition, int]
+        ] = defaultdict(dict)
 
         # The active subscriptions are stored by consumer group as a mapping
         # between the consumer and it's subscribed topics.
@@ -326,7 +326,7 @@ class LocalConsumer(Consumer[TStrategyPayload]):
             # atomic
             self.__staged_offsets.update(offsets)
 
-    def commit_offsets(self) -> Mapping[Partition, int]:
+    def commit_offsets(self) -> Optional[Mapping[Partition, int]]:
         with self.__lock:
             if self.__closed:
                 raise RuntimeError("consumer is closed")

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/arroyo/processing/processor.py

@@ -146,9 +146,9 @@ class StreamProcessor(Generic[TStrategyPayload]):
         self.__processor_factory = processor_factory
         self.__metrics_buffer = MetricsBuffer()
 
-        self.__processing_strategy: Optional[
-            ProcessingStrategy[TStrategyPayload]
-        ] = None
+        self.__processing_strategy: Optional[ProcessingStrategy[TStrategyPayload]] = (
+            None
+        )
 
         self.__message: Optional[BrokerValue[TStrategyPayload]] = None
 
@@ -465,8 +465,9 @@ class StreamProcessor(Generic[TStrategyPayload]):
 
                     elif self.__is_paused:
                         paused_partitions = set(self.__consumer.paused())
+                        all_partitions = set(self.__consumer.tell())
                         unpaused_partitions = (
-                            set(self.__consumer.tell()) - paused_partitions
+                            all_partitions - paused_partitions
                         )
                         if unpaused_partitions:
                             logger.warning(
@@ -484,6 +485,18 @@ class StreamProcessor(Generic[TStrategyPayload]):
                             # A paused consumer should still poll periodically to avoid it's partitions
                             # getting revoked by the broker after reaching the max.poll.interval.ms
                             # Polling a paused consumer should never yield a message.
+                            logger.warning("consumer.tell() value right before poll() is: %s", self.__consumer.tell())
+                            maybe_message = self.__consumer.poll(0.1)
+                            if maybe_message is not None:
+                                logger.warning("Received a message from partition: %s, \
+                                                consumer.tell() value right after poll() is: %s \
+                                                Some lines above consumer.tell() was called, all_partitons value was: %s \
+                                                Some lines above consumer.paused() was called, paused_partitions value is: %s",
+                                                maybe_message.partition,
+                                                self.__consumer.tell(),
+                                                all_partitions,
+                                                paused_partitions
+                                                )
                             assert self.__consumer.poll(0.1) is None
                     else:
                         time.sleep(0.01)

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0/sentry_arroyo.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.30.0
+Version: 2.32.0
 Summary: Arroyo is a Python library for working with streaming data.
 Home-page: https://github.com/getsentry/arroyo
 Author: Sentry

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/sentry_arroyo.egg-info/SOURCES.txt

@@ -64,6 +64,7 @@ tests/test_types.py
 tests/backends/__init__.py
 tests/backends/mixins.py
 tests/backends/test_commit.py
+tests/backends/test_confluent_producer.py
 tests/backends/test_kafka.py
 tests/backends/test_kafka_producer.py
 tests/backends/test_local.py

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/setup.py

@@ -10,7 +10,7 @@ def get_requirements() -> Sequence[str]:
 
 setup(
     name="sentry-arroyo",
-    version="2.30.0",
+    version="2.32.0",
     author="Sentry",
     author_email="oss@sentry.io",
     license="Apache-2.0",

sentry_arroyo-2.32.0/tests/backends/test_confluent_producer.py

@@ -0,0 +1,79 @@
+from typing import Optional
+from unittest import mock
+
+from confluent_kafka import KafkaError
+from confluent_kafka import Message as ConfluentMessage
+from confluent_kafka import Producer as ConfluentKafkaProducer
+
+from arroyo.backends.kafka.consumer import ConfluentProducer
+from tests.metrics import Increment, TestingMetricsBackend
+
+
+class TestConfluentProducer:
+    """
+    Tests for ConfluentProducer wrapper around confluent_kafka.Producer.
+    """
+
+    def test_init(self) -> None:
+        """Test that ConfluentProducer can be instantiated"""
+        config = {"bootstrap.servers": "fake:9092"}
+        producer = ConfluentProducer(config)
+
+        assert isinstance(producer, ConfluentProducer)
+        assert isinstance(producer, ConfluentKafkaProducer)
+
+    def test_metrics_callback_records_success(self) -> None:
+        """Test that the metrics callback records success metric"""
+        producer = ConfluentProducer(
+            {"bootstrap.servers": "fake:9092", "client.id": "test-producer-name"}
+        )
+        mock_message = mock.Mock(spec=ConfluentMessage)
+        producer._ConfluentProducer__metrics_delivery_callback(None, mock_message)
+        producer.flush()  # Flush buffered metrics
+        assert (
+            Increment(
+                "arroyo.producer.produce_status",
+                1,
+                {"status": "success", "producer_name": "test-producer-name"},
+            )
+            in TestingMetricsBackend.calls
+        )
+
+    def test_metrics_callback_records_error(self) -> None:
+        """Test that the metrics callback records error metric"""
+        producer = ConfluentProducer({"bootstrap.servers": "fake:9092"})
+        mock_error = mock.Mock(spec=KafkaError)
+        mock_message = mock.Mock(spec=ConfluentMessage)
+        producer._ConfluentProducer__metrics_delivery_callback(mock_error, mock_message)
+        producer.flush()  # Flush buffered metrics
+        assert (
+            Increment("arroyo.producer.produce_status", 1, {"status": "error"})
+            in TestingMetricsBackend.calls
+        )
+
+    def test_delivery_callback_wraps_user_callback(self) -> None:
+        """Test that the delivery callback wrapper calls both metrics and user callbacks"""
+        producer = ConfluentProducer(
+            {"bootstrap.servers": "fake:9092", "client.id": "test-producer-name"}
+        )
+        user_callback_invoked = []
+
+        def user_callback(
+            error: Optional[KafkaError], message: ConfluentMessage
+        ) -> None:
+            user_callback_invoked.append((error, message))
+
+        wrapped = producer._ConfluentProducer__delivery_callback(user_callback)
+        mock_message = mock.Mock(spec=ConfluentMessage)
+        wrapped(None, mock_message)
+        producer.flush()  # Flush buffered metrics
+        assert (
+            Increment(
+                "arroyo.producer.produce_status",
+                1,
+                {"status": "success", "producer_name": "test-producer-name"},
+            )
+            in TestingMetricsBackend.calls
+        )
+        assert len(user_callback_invoked) == 1
+        assert user_callback_invoked[0] == (None, mock_message)

{sentry_arroyo-2.30.0 → sentry_arroyo-2.32.0}/tests/backends/test_kafka.py

@@ -275,6 +275,70 @@ class TestKafkaStreams(StreamsTestMixin[KafkaPayload]):
                 processor._run_once()
                 assert consumer.paused() == []
 
+    def test_auto_commit_mode(self) -> None:
+        """Test that auto-commit mode uses store_offsets and commits on close"""
+        group_id = uuid.uuid1().hex
+
+        with self.get_topic() as topic:
+            # Produce some messages
+            with closing(self.get_producer()) as producer:
+                for i in range(5):
+                    payload = KafkaPayload(None, f"msg_{i}".encode("utf8"), [])
+                    producer.produce(topic, payload).result(5.0)
+
+            # Create consumer with auto-commit enabled
+            configuration = {
+                **self.configuration,
+                "auto.offset.reset": "earliest",
+                "arroyo.enable.auto.commit": True,
+                "group.id": group_id,
+                "session.timeout.ms": 10000,
+            }
+
+            # First consumer: consume messages and close
+            consumed_offsets = []
+            with closing(KafkaConsumer(configuration)) as consumer:
+                consumer.subscribe([topic])
+
+                # Consume all 5 messages and stage their offsets
+                for i in range(5):
+                    value = consumer.poll(10.0)
+                    assert value is not None
+                    consumed_offsets.append(value.offset)
+
+                    # Stage offsets (will use store_offsets internally in auto-commit mode)
+                    consumer.stage_offsets(value.committable)
+
+                # commit_offsets should return None in auto-commit mode
+                result = consumer.commit_offsets()
+                assert result is None
+
+                # Close will commit any stored offsets
+
+            # Verify we consumed offsets 0-4
+            assert consumed_offsets == [0, 1, 2, 3, 4]
+
+            # Second consumer: verify offsets were committed on close
+            # This consumer uses manual commit to verify the committed offset
+            with closing(
+                self.get_consumer(
+                    group=group_id,
+                    auto_offset_reset="earliest",
+                    enable_end_of_partition=True,
+                )
+            ) as consumer:
+                consumer.subscribe([topic])
+
+                # Should start from offset 5, hitting EndOfPartition immediately
+                # If we got a message with offset < 5, auto-commit didn't work
+                try:
+                    consumer.poll(10.0)
+                    pytest.fail("Expected EndOfPartition, but poll succeeded")
+                except EndOfPartition as e:
+                    # Verify we got EndOfPartition at offset 5
+                    assert e.offset == 5
+                    assert e.partition == Partition(topic, 0)
+
 
 class TestKafkaStreamsIncrementalRebalancing(TestKafkaStreams):
     # re-test the kafka consumer with cooperative-sticky rebalancing