Adding support for pipeline DLQ to SQS sink

divbok · divbok · commit cb5347ac725a · 2026-03-03T15:39:35.000-06:00
Signed-off-by: Divyansh Bokadia &lt;dbokadia@amazon.com&gt;
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventFailureMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventFailureMetadata.java
@@ -11,5 +11,15 @@
 
 public interface EventFailureMetadata {
     EventFailureMetadata with(String key, Object value);
+
+    EventFailureMetadata withPluginId(String value);
+
+    EventFailureMetadata withPluginName(String value);
+
+    EventFailureMetadata withPipelineName(String value);
+
+    EventFailureMetadata withErrorMessage(Object value);
+
+
 }
 
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
@@ -64,11 +64,32 @@
 public class JacksonEvent implements Event {
     class DefaultEventFailureMetadata implements EventFailureMetadata {
         static final String FAILURE_METADATA = "_failure_metadata";
+        static final String PLUGIN_ID = "pluginId";
+        static final String PLUGIN_NAME = "pluginName";
+        static final String PIPELINE_NAME = "pipelineName";
+        static final String ERROR_MESSAGE = "errorMessage";
+
 
         public DefaultEventFailureMetadata with(String key, Object value) {
             put(FAILURE_METADATA+"/"+key, value);
             return this;
         }
+
+        public DefaultEventFailureMetadata withPluginId(String value) {
+            return with(PLUGIN_ID, value);
+        }
+
+        public DefaultEventFailureMetadata withPluginName(String value) {
+            return with(PLUGIN_NAME, value);
+        }
+
+        public DefaultEventFailureMetadata withPipelineName(String value) {
+            return with(PIPELINE_NAME, value);
+        }
+
+        public DefaultEventFailureMetadata withErrorMessage(Object value) {
+            return with(ERROR_MESSAGE, value);
+        }
     }
 
     private static final Logger LOG = LoggerFactory.getLogger(JacksonEvent.class);
diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java
@@ -202,15 +202,28 @@ void testUpdateFailureMetadata() {
     @Test
     public void testDefaultEventFailureMetadata() {
         String eventType = UUID.randomUUID().toString();
-
+        String pluginId = UUID.randomUUID().toString();
+        String pluginName = UUID.randomUUID().toString();
+        String pipelineName = UUID.randomUUID().toString();
+        String errorMessage = UUID.randomUUID().toString();
         Event event = JacksonEvent.builder()
                 .withEventType(eventType)
                 .build();
 
         EventFailureMetadata eventFailureMetadata = event.updateFailureMetadata();
         eventFailureMetadata.with("key1", "value1").with("key2", 2);
-        assertThat(event.get(JacksonEvent.DefaultEventFailureMetadata.FAILURE_METADATA+"/key1", String.class), equalTo("value1"));
-        assertThat(event.get(JacksonEvent.DefaultEventFailureMetadata.FAILURE_METADATA+"/key2", Integer.class), equalTo(2));
+        eventFailureMetadata.withPluginId(pluginId);
+        eventFailureMetadata.withPluginName(pluginName);
+        eventFailureMetadata.withPipelineName(pipelineName);
+        eventFailureMetadata.withErrorMessage(errorMessage);
+        String base = JacksonEvent.DefaultEventFailureMetadata.FAILURE_METADATA + "/";
+        assertThat(event.get(base + "key1", String.class), equalTo("value1"));
+        assertThat(event.get(base + "key2", Integer.class), equalTo(2));
+        assertThat(event.get(base + JacksonEvent.DefaultEventFailureMetadata.PLUGIN_ID, String.class), equalTo(pluginId));
+        assertThat(event.get(base + JacksonEvent.DefaultEventFailureMetadata.PLUGIN_NAME, String.class), equalTo(pluginName));
+        assertThat(event.get(base + JacksonEvent.DefaultEventFailureMetadata.PIPELINE_NAME, String.class), equalTo(pipelineName));
+        assertThat(event.get(base + JacksonEvent.DefaultEventFailureMetadata.ERROR_MESSAGE, String.class), equalTo(errorMessage));
+
     }
 
     @Test
diff --git a/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java b/data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java
@@ -641,10 +641,10 @@ private void logFailureForDlqObjects(final List<DlqObject> dlqObjects, final Thr
 
         if (event != null) {
             event.updateFailureMetadata()
-                .with("pluginId", dlqObject.getPluginId())
-                .with("pluginName", dlqObject.getPluginName())
-                .with("pipelineName", dlqObject.getPipelineName())
-                .with("message",  ((FailedDlqData) dlqObject.getFailedData()).getMessage())
+                .withPluginId(dlqObject.getPluginId())
+                .withPluginName(dlqObject.getPluginName())
+                .withPipelineName(dlqObject.getPipelineName())
+                .withErrorMessage(((FailedDlqData) dlqObject.getFailedData()).getMessage())
                 .with("status", ((FailedDlqData) dlqObject.getFailedData()).getStatus())
                 .with("index", ((FailedDlqData) dlqObject.getFailedData()).getIndex())
                 .with("indexId", ((FailedDlqData) dlqObject.getFailedData()).getIndexId());
diff --git a/data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/service/PrometheusSinkService.java b/data-prepper-plugins/prometheus-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/prometheus/service/PrometheusSinkService.java
@@ -93,11 +93,11 @@ public void addFailedEventsToDlq(final List<Event> failedEvents, final Throwable
             }
             event.updateFailureMetadata()
                 .with("statusCode", statusCode)
-                .with("pluginName", PLUGIN_NAME)
-                .with("pipelineName", pipelineDescription.getPipelineName());
+                .withPluginName(PLUGIN_NAME)
+                .withPipelineName(pipelineDescription.getPipelineName());
             if (ex != null) {
                 event.updateFailureMetadata()
-                    .with("message",  ex.getMessage());
+                    .withErrorMessage(ex.getMessage());
             }
             dlqRecords.add(new Record<>(event));
         }
diff --git a/data-prepper-plugins/sqs-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sqs/SqsSinkService.java b/data-prepper-plugins/sqs-sink/src/main/java/org/opensearch/dataprepper/plugins/sink/sqs/SqsSinkService.java
@@ -316,11 +316,11 @@ private void addMessageToDLQ(final String message, final List<EventHandle> event
             for (Event event: events) {
                 if (event != null) {
                    event.updateFailureMetadata()
-                       .with("pluginId", pluginSetting.getName())
-                       .with("pluginName", pluginSetting.getName())
-                       .with("pipelineName", pipelineDescription.getPipelineName())
-                       .with("sqsSinkQueueUrl", queueUrl)
-                       .with("message", errorMessage);
+                       .withPluginId(pluginSetting.getName())
+                       .withPluginName(pluginSetting.getName())
+                       .withPipelineName(pipelineDescription.getPipelineName())
+                       .withErrorMessage(errorMessage)
+                       .with("sqsSinkQueueUrl", queueUrl);
                     dlqPipelineRecords.add(new Record<>(event));
                 }
             }
diff --git a/data-prepper-plugins/sqs-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/sqs/SqsSinkServiceTest.java b/data-prepper-plugins/sqs-sink/src/test/java/org/opensearch/dataprepper/plugins/sink/sqs/SqsSinkServiceTest.java
@@ -128,6 +128,9 @@ class SqsSinkServiceTest {
     private OutputCodecContext outputCodecContext;
     private String queueUrl;
 
+    private String pluginName;
+    private String pipelineName;
+
     private SqsSinkService createObjectUnderTest() {
         return new SqsSinkService(sqsSinkConfig, sqsClient, expressionEvaluator, outputCodec, sinkContext, dlqPushHandler,
                 pluginMetrics, pluginSetting, pipelineDescription);
@@ -151,14 +154,16 @@ void setup() {
         when (thresholdConfig.getMaxMessageSizeBytes()).thenReturn(256*1024L);
         when (thresholdConfig.getMaxEventsPerMessage()).thenReturn(1);
         when (sqsSinkConfig.getThresholdConfig()).thenReturn(thresholdConfig);
-        when (pipelineDescription.getPipelineName()).thenReturn("pipeline");
+        pipelineName = UUID.randomUUID().toString();
+        when (pipelineDescription.getPipelineName()).thenReturn(pipelineName);
         lenient().when (sqsSinkConfig.getMaxRetries()).thenReturn(3);
         lenient().when(flushResponse.hasFailed()).thenReturn(false);
         when(sqsClient.sendMessageBatch(any(SendMessageBatchRequest.class))).thenReturn(flushResponse);
         when(expressionEvaluator.isValidFormatExpression(anyString())).thenReturn(true);
         when(dlqPushHandler.perform(any(List.class))).thenReturn(true);
-        lenient().when(pluginSetting.getName()).thenReturn("name");
-        lenient().when(pluginSetting.getPipelineName()).thenReturn("pipeline");
+        pluginName = UUID.randomUUID().toString();
+        lenient().when(pluginSetting.getName()).thenReturn(pluginName);
+        lenient().when(pluginSetting.getPipelineName()).thenReturn(pipelineName);
         when(dlqPushHandler.getPluginSetting()).thenReturn(pluginSetting);
         lenient().doNothing().when(summary).record(any(Double.class));
         lenient().doNothing().when(timer).record(any(Long.class), any(TimeUnit.class));
@@ -535,7 +540,6 @@ void TestWithOneBatch_RetryFlushes() throws Exception {
     @Test
     void TestLargeRecordToPipelineDLQ() {
         dlqPipeline = mock(HeadlessPipeline.class);
-
         List<Record<Event>> capturedRecords = new ArrayList<>();
         doAnswer(invocation -> {
             List<Record<Event>> arg = invocation.getArgument(0);
@@ -551,9 +555,9 @@ void TestLargeRecordToPipelineDLQ() {
         assertThat(capturedRecords.size(), equalTo(1));
 
         Event event = capturedRecords.get(0).getData();
-        assertThat(event.get("_failure_metadata/pluginId", String.class), equalTo("name"));
-        assertThat(event.get("_failure_metadata/pluginName", String.class), equalTo("name"));
-        assertThat(event.get("_failure_metadata/pipelineName", String.class), equalTo("pipeline"));
+        assertThat(event.get("_failure_metadata/pluginId", String.class), equalTo(pluginName));
+        assertThat(event.get("_failure_metadata/pluginName", String.class), equalTo(pluginName));
+        assertThat(event.get("_failure_metadata/pipelineName", String.class), equalTo(pipelineName));
         assertThat(event.get("_failure_metadata/sqsSinkQueueUrl", String.class), equalTo(queueUrl));
         verify(eventHandle, never()).release(anyBoolean());
     }

Original file line number	Diff line number	Diff line change
`@@ -93,11 +93,11 @@ public void addFailedEventsToDlq(final List<Event> failedEvents, final Throwable`
`93`	`93`	`}`
`94`	`94`	`event.updateFailureMetadata()`
`95`	`95`	`.with("statusCode", statusCode)`
`96`		`- .with("pluginName", PLUGIN_NAME)`
`97`		`- .with("pipelineName", pipelineDescription.getPipelineName());`
	`96`	`+ .withPluginName(PLUGIN_NAME)`
	`97`	`+ .withPipelineName(pipelineDescription.getPipelineName());`
`98`	`98`	`if (ex != null) {`
`99`	`99`	`event.updateFailureMetadata()`
`100`		`- .with("message", ex.getMessage());`
	`100`	`+ .withErrorMessage(ex.getMessage());`
`101`	`101`	`}`
`102`	`102`	`dlqRecords.add(new Record<>(event));`
`103`	`103`	`}`