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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.30.0
+Version: 2.31.1
 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.31.1}/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.31.1}/arroyo/backends/kafka/consumer.py

@@ -33,7 +33,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 +50,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__)
@@ -657,7 +658,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 +743,57 @@ 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]
+
+
+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.__metrics = get_metrics()
+        self.producer_name = configuration.get("client.id") or None
+
+    def __metrics_delivery_callback(
+        self,
+        error: Optional[KafkaError],
+        _message: ConfluentMessage,
+    ) -> None:
+        if error is not None:
+            status = "error"
+        else:
+            status = "success"
+
+        tags = {"status": status}
+        if self.producer_name:
+            tags["producer_name"] = self.producer_name
+
+        self.__metrics.increment(
+            "arroyo.producer.produce_status",
+            1,
+            tags=tags,
+        )
+
+    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)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sentry-arroyo
-Version: 2.30.0
+Version: 2.31.1
 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.31.1}/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.31.1}/setup.py

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

sentry_arroyo-2.31.1/tests/backends/test_confluent_producer.py

@@ -0,0 +1,76 @@
+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)
+        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)
+        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)
+        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)