Preserve DLQ producer override_params from config

Author: cursoragent

sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py

@@ -186,7 +186,7 @@ def build_dlq_config(
         topic=step.dlq_config.topic,
         producer_config=PyKafkaProducerConfig(
             bootstrap_servers=step.dlq_config.bootstrap_servers,
-            override_params=None,
+            override_params=step.dlq_config.override_params,
         ),
     )
 

sentry_streams/sentry_streams/pipeline/pipeline.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import timedelta
 from enum import Enum
 from functools import partial
@@ -235,6 +235,7 @@ class DlqConfig:
 
     topic: str
     bootstrap_servers: Sequence[str]
+    override_params: Mapping[str, str] = field(default_factory=dict)
 
 
 @dataclass
@@ -267,13 +268,17 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None:
             servers = producer_config.get("bootstrap_servers") or (
                 self.dlq_config.bootstrap_servers if self.dlq_config else None
             )
+            override_params = producer_config.get("override_params")
+            if override_params is None:
+                override_params = self.dlq_config.override_params if self.dlq_config else {}
 
             if topic is None or servers is None:
                 raise ValueError("DLQ config requires both 'topic' and 'bootstrap_servers'")
 
             self.dlq_config = DlqConfig(
                 topic=topic,
                 bootstrap_servers=servers,
+                override_params=cast(Mapping[str, str], override_params),
             )
 
 

sentry_streams/tests/test_dlq.py

@@ -68,21 +68,28 @@ def test_consumer_creation(
 
 
 @pytest.mark.parametrize(
-    "dlq_config, expected_topic, expected_bootstrap_servers",
+    "dlq_config, expected_topic, expected_bootstrap_servers, expected_override_params",
     [
-        pytest.param(None, None, None, id="no_dlq_config"),
+        pytest.param(None, None, None, None, id="no_dlq_config"),
         pytest.param(
-            DlqConfig(topic="test-dlq", bootstrap_servers=["localhost:9092"]),
+            DlqConfig(
+                topic="test-dlq",
+                bootstrap_servers=["localhost:9092"],
+                override_params={"sasl.username": "test"},
+            ),
             "test-dlq",
             ["localhost:9092"],
+            {"sasl.username": "test"},
             id="single_bootstrap_server",
         ),
         pytest.param(
             DlqConfig(
-                topic="my-dlq", bootstrap_servers=["broker1:9092", "broker2:9092", "broker3:9092"]
+                topic="my-dlq",
+                bootstrap_servers=["broker1:9092", "broker2:9092", "broker3:9092"],
             ),
             "my-dlq",
             ["broker1:9092", "broker2:9092", "broker3:9092"],
+            {},
             id="multiple_bootstrap_servers",
         ),
     ],
@@ -91,6 +98,7 @@ def test_build_dlq_config(
     dlq_config: DlqConfig | None,
     expected_topic: str | None,
     expected_bootstrap_servers: list[str] | None,
+    expected_override_params: dict[str, str] | None,
 ) -> None:
     """Test build_dlq_config returns correct PyDlqConfig for various inputs."""
     source = StreamSource(
@@ -109,47 +117,75 @@ def test_build_dlq_config(
         assert result.topic == expected_topic
         assert result.producer_config is not None
         assert result.producer_config.bootstrap_servers == expected_bootstrap_servers
-        assert result.producer_config.override_params is None
+        assert result.producer_config.override_params == expected_override_params
 
 
 @pytest.mark.parametrize(
-    "initial_dlq_config, override_dlq, expected_topic, expected_bootstrap_servers",
+    "initial_dlq_config, override_dlq, expected_topic, expected_bootstrap_servers, expected_override_params",
     [
         pytest.param(
             None,
-            {"topic": "new-dlq", "producer_config": {"bootstrap_servers": ["broker1:9092"]}},
+            {
+                "topic": "new-dlq",
+                "producer_config": {
+                    "bootstrap_servers": ["broker1:9092"],
+                    "override_params": {"security.protocol": "SASL_SSL"},
+                },
+            },
             "new-dlq",
             ["broker1:9092"],
+            {"security.protocol": "SASL_SSL"},
             id="create_new_config",
         ),
         pytest.param(
-            DlqConfig(topic="old-dlq", bootstrap_servers=["old-broker:9092"]),
+            DlqConfig(
+                topic="old-dlq",
+                bootstrap_servers=["old-broker:9092"],
+                override_params={"old.param": "old-value"},
+            ),
             {"topic": "new-dlq"},
             "new-dlq",
             ["old-broker:9092"],
+            {"old.param": "old-value"},
             id="override_topic_only",
         ),
         pytest.param(
-            DlqConfig(topic="old-dlq", bootstrap_servers=["old-broker:9092"]),
-            {"producer_config": {"bootstrap_servers": ["new-broker:9092", "new-broker2:9092"]}},
+            DlqConfig(
+                topic="old-dlq",
+                bootstrap_servers=["old-broker:9092"],
+                override_params={"old.param": "old-value"},
+            ),
+            {
+                "producer_config": {
+                    "bootstrap_servers": ["new-broker:9092", "new-broker2:9092"],
+                    "override_params": {"new.param": "new-value"},
+                }
+            },
             "old-dlq",
             ["new-broker:9092", "new-broker2:9092"],
+            {"new.param": "new-value"},
             id="override_bootstrap_servers_only",
         ),
         pytest.param(
-            DlqConfig(topic="old-dlq", bootstrap_servers=["old-broker:9092"]),
+            DlqConfig(
+                topic="old-dlq",
+                bootstrap_servers=["old-broker:9092"],
+                override_params={"old.param": "old-value"},
+            ),
             {"topic": "new-dlq", "producer_config": {"bootstrap_servers": ["new-broker:9092"]}},
             "new-dlq",
             ["new-broker:9092"],
+            {"old.param": "old-value"},
             id="override_both_fields",
         ),
     ],
 )
 def test_stream_source_override_config_dlq(
     initial_dlq_config: DlqConfig | None,
-    override_dlq: dict[str, str | list[str]],
+    override_dlq: dict[str, object],
     expected_topic: str,
     expected_bootstrap_servers: list[str],
+    expected_override_params: dict[str, str],
 ) -> None:
     """Test that StreamSource.override_config correctly overrides DLQ settings."""
     source = StreamSource(
@@ -163,3 +199,4 @@ def test_stream_source_override_config_dlq(
     assert source.dlq_config is not None
     assert source.dlq_config.topic == expected_topic
     assert source.dlq_config.bootstrap_servers == expected_bootstrap_servers
+    assert source.dlq_config.override_params == expected_override_params