Fix dlqPipeline functionality broken by PR 6349 (#6678)

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

Author: kkondaka

data-prepper-plugins/opensearch/src/integrationTest/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSinkIT.java

@@ -71,6 +71,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Collection;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
@@ -103,6 +104,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import org.mockito.ArgumentCaptor;
 import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchIntegrationHelper.createContentParser;
 import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchIntegrationHelper.createOpenSearchClient;
 import static org.opensearch.dataprepper.plugins.sink.opensearch.OpenSearchIntegrationHelper.getHosts;
@@ -173,6 +175,23 @@ private OpenSearchSink createObjectUnderTest(OpenSearchSinkConfig openSearchSink
         return sink;
     }
 
+    private OpenSearchSink createObjectUnderTestWithDlqPipeline(OpenSearchSinkConfig openSearchSinkConfig, HeadlessPipeline dlqPipeline, boolean doInitialize) {
+        sinkContext = mock(SinkContext.class);
+        when(sinkContext.getTagsTargetKey()).thenReturn(null);
+        when(sinkContext.getForwardToPipelines()).thenReturn(Map.of());
+        when(pipelineDescription.getPipelineName()).thenReturn(PIPELINE_NAME);
+        when(pluginSetting.getPipelineName()).thenReturn(PIPELINE_NAME);
+        when(pluginSetting.getName()).thenReturn(PLUGIN_NAME);
+        OpenSearchSink sink = new OpenSearchSink(
+                pluginSetting, sinkContext, expressionEvaluator, awsCredentialsSupplier, pipelineDescription, pluginConfigObservable, openSearchSinkConfig);
+        sink.setFailurePipeline(dlqPipeline);
+        if (doInitialize) {
+            sink.doInitialize();
+        }
+        sinksToShutdown.add(sink);
+        return sink;
+    }
+
     private OpenSearchSink createObjectUnderTestWithSinkContext(OpenSearchSinkConfig openSearchSinkConfig, final Map<String, HeadlessPipeline> forwardPipelineMap, boolean doInitialize) {
         sinkContext = mock(SinkContext.class);
         testTagsTargetKey = RandomStringUtils.randomAlphabetic(5);
@@ -906,6 +925,47 @@ void testOutputForwardsCreatedDocumentsToAPipeline() throws IOException, Interru
         verify(sinkContext).forwardRecords(any(), eq(null), eq(null));
     }
 
+    @Test
+    @Timeout(value = 50, unit = TimeUnit.SECONDS)
+    void testOutputFailedDocumentsToDLQPipeline() throws IOException, InterruptedException {
+        HeadlessPipeline dlqPipeline = mock(HeadlessPipeline.class);
+        final String testIndexAlias = "test-alias";
+        final String testTemplateFile = Objects.requireNonNull(
+                getClass().getClassLoader().getResource(TEST_TEMPLATE_V1_FILE)).getFile();
+        final String testIdField = "someId";
+        final String testId = "foo";
+        final String testId2 = "foo2";
+        final Record<Event> testRecord1 = new Record(JacksonEvent.builder()
+                    .withEventType(EventType.TRACE.toString())
+                    .withData(Map.of(testIdField, testId, "name", "value")).build());
+        final Record<Event> testRecord2 = new Record(JacksonEvent.builder()
+                    .withEventType(EventType.TRACE.toString())
+                    .withData(Map.of(testIdField, testId, "name", Map.of("key", "value"))).build());
+        //List<Record<Event>> testRecords = Arrays.asList(testRecord1, testRecord2);
+        List<Record<Event>> testRecords = Collections.singletonList(testRecord1);
+        Map<String, Object> metadata = initializeConfigurationMetadata(null, testIndexAlias, testTemplateFile);
+        metadata.put(IndexConfiguration.DOCUMENT_ID_FIELD, testIdField);
+        final OpenSearchSinkConfig openSearchSinkConfig = generateOpenSearchSinkConfigByMetadata(metadata);
+        final OpenSearchSink sink = createObjectUnderTestWithDlqPipeline(openSearchSinkConfig, dlqPipeline, true);
+        sink.output(testRecords);
+        final List<Map<String, Object>> retSources = getSearchResponseDocSources(testIndexAlias);
+        assertThat(retSources.size(), equalTo(1));
+        assertThat(getDocumentCount(testIndexAlias, "_id", testId), equalTo(Integer.valueOf(1)));
+
+        // verify metrics
+        final List<Measurement> bulkRequestLatencies = MetricsTestUtil.getMeasurementList(
+                new StringJoiner(MetricNames.DELIMITER).add(PIPELINE_NAME).add(PLUGIN_NAME)
+                        .add(OpenSearchSink.BULKREQUEST_LATENCY).toString());
+        assertThat(bulkRequestLatencies.size(), equalTo(3));
+        // COUNT
+        Assert.assertEquals(1.0, bulkRequestLatencies.get(0).getValue(), 0);
+        testRecords = Collections.singletonList(testRecord2);
+        sink.output(testRecords);
+        ArgumentCaptor<Collection<Record<Event>>> captor = ArgumentCaptor.forClass(Collection.class);
+        verify(dlqPipeline).sendEvents(captor.capture());
+        assertThat(captor.getValue().size(), equalTo(1));
+    }
+
     @Test
     @Timeout(value = 50, unit = TimeUnit.SECONDS)
     void testOutputCustomIndex() throws IOException, InterruptedException {

data-prepper-plugins/opensearch/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/OpenSearchSink.java

@@ -157,7 +157,7 @@ public class OpenSearchSink extends AbstractSink<Record<Event>> {
   private volatile boolean initialized;
   private final SinkContext sinkContext;
   private final ExpressionEvaluator expressionEvaluator;
-  private final boolean useEventInBulkOperation;
+  private boolean useEventInBulkOperation;
 
   private FailedBulkOperationConverter failedBulkOperationConverter;
   private DataStreamDetector dataStreamDetector;
@@ -191,7 +191,6 @@ public OpenSearchSink(final PluginSetting pluginSetting,
     sinkForwardRecordsContext = new SinkForwardRecordsContext(sinkContext);
     this.expressionEvaluator = expressionEvaluator;
     this.pipeline = pipelineDescription.getPipelineName();
-    this.useEventInBulkOperation = (getFailurePipeline() != null || sinkContext.getForwardToPipelines().size() > 0);
     bulkRequestTimer = pluginMetrics.timer(BULKREQUEST_LATENCY);
     bulkRequestErrorsCounter = pluginMetrics.counter(BULKREQUEST_ERRORS);
     invalidActionErrorsCounter = pluginMetrics.counter(INVALID_ACTION_ERRORS);
@@ -234,6 +233,9 @@ public OpenSearchSink(final PluginSetting pluginSetting,
   public void doInitialize() {
     try {
         doInitializeInternal();
+        // getFailurePipeline() does not return valid value in the constructor. Earliest it can be used
+        // is in doInitialize()
+        useEventInBulkOperation = (getFailurePipeline() != null || sinkContext.getForwardToPipelines().size() > 0);
     } catch (IOException e) {
         LOG.warn("Failed to initialize OpenSearch sink, retrying: {} ", e.getMessage());
         this.shutdown();