Wire together for opensearch findings

Signed-off-by: Chase Engelbrecht <[email protected]>
engechas · Mar 16, 2024 · 48df422 · 48df422
1 parent 4656adc
commit 48df422
Show file tree

Hide file tree

Showing 17 changed files with 544 additions and 180 deletions.
diff --git a/...rc/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java b/...rc/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkOperationWrapper.java
@@ -6,12 +6,14 @@
 package org.opensearch.dataprepper.plugins.sink.opensearch;
 
 import org.opensearch.client.opensearch.core.bulk.BulkOperation;
+import org.opensearch.client.opensearch.core.bulk.BulkResponseItem;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.event.EventHandle;
 import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.SerializedJson;
 
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -44,24 +46,31 @@ public class BulkOperationWrapper {
     );
 
     private final EventHandle eventHandle;
+    private final Consumer<BulkResponseItem> bulkResponseItemConsumer;
     private final BulkOperation bulkOperation;
     private final SerializedJson jsonNode;
     private final Event event;
 
     public BulkOperationWrapper(final BulkOperation bulkOperation) {
-        this(bulkOperation, null, null, null);
+        this(bulkOperation, null, null, null, null);
     }
 
-    public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle, final SerializedJson jsonNode, final Event event) {
+    public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle, final SerializedJson jsonNode, final Event event,
+                                final Consumer<BulkResponseItem> bulkResponseItemConsumer) {
         checkNotNull(bulkOperation);
         this.bulkOperation = bulkOperation;
         this.eventHandle = eventHandle;
         this.jsonNode = jsonNode;
         this.event = event;
+        this.bulkResponseItemConsumer = bulkResponseItemConsumer;
+    }
+
+    public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle, final SerializedJson jsonNode, final Event event) {
+        this(bulkOperation, eventHandle, jsonNode, event, null);
     }
 
     public BulkOperationWrapper(final BulkOperation bulkOperation, final EventHandle eventHandle) {
-        this(bulkOperation, eventHandle, null, null);
+        this(bulkOperation, eventHandle, null, null, null);
     }
 
     public BulkOperation getBulkOperation() {
@@ -95,6 +104,10 @@ public String getId() {
         return getValueFromConverter(BULK_OPERATION_TO_ID_CONVERTERS);
     }
 
+    public Consumer<BulkResponseItem> getBulkResponseItemConsumer() {
+        return bulkResponseItemConsumer;
+    }
+
     private <T> T getValueFromConverter(final Map<Predicate<BulkOperation>, Function<BulkOperation, T>> converters) {
         final List<T> values = converters.entrySet().stream()
                 .filter(entry -> entry.getKey().test(bulkOperation))

diff --git a/...h/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java b/...h/src/main/java/org/opensearch/dataprepper/plugins/sink/opensearch/BulkRetryStrategy.java
@@ -301,6 +301,7 @@ private BulkOperationRequestResponse handleRetry(final AccumulatingBulkRequest r
                 final BulkResponseItem bulkResponseItem = bulkResponse.items().get(i);
 
                 bulkOperation.releaseEventHandle(true);
+                executeConsumer(bulkOperation, bulkResponseItem);
             }
         }
         return null;
@@ -325,6 +326,7 @@ private AccumulatingBulkRequest<BulkOperationWrapper, BulkRequest> createBulkReq
                         documentsVersionConflictErrors.increment();
                         LOG.debug("Received version conflict from OpenSearch: {}", bulkItemResponse.error().reason());
                         bulkOperation.releaseEventHandle(true);
+                        executeConsumer(bulkOperation, bulkItemResponse);
                     } else {
                         nonRetryableFailures.add(FailedBulkOperation.builder()
                                 .withBulkOperation(bulkOperation)
@@ -335,6 +337,7 @@ private AccumulatingBulkRequest<BulkOperationWrapper, BulkRequest> createBulkReq
                 } else {
                     sentDocumentsCounter.increment();
                     bulkOperation.releaseEventHandle(true);
+                    executeConsumer(bulkOperation, bulkItemResponse);
                 }
                 index++;
             }
@@ -357,6 +360,7 @@ private void handleFailures(final AccumulatingBulkRequest<BulkOperationWrapper,
                     documentsVersionConflictErrors.increment();
                     LOG.debug("Received version conflict from OpenSearch: {}", bulkItemResponse.error().reason());
                     bulkOperation.releaseEventHandle(true);
+                    executeConsumer(bulkOperation, bulkItemResponse);
                 } else {
                     failures.add(FailedBulkOperation.builder()
                             .withBulkOperation(bulkOperation)
@@ -367,6 +371,7 @@ private void handleFailures(final AccumulatingBulkRequest<BulkOperationWrapper,
             } else {
                 sentDocumentsCounter.increment();
                 bulkOperation.releaseEventHandle(true);
+                executeConsumer(bulkOperation, bulkItemResponse);
             }
         }
         logFailure.accept(failures.build(), null);
@@ -387,4 +392,12 @@ private void handleFailures(final AccumulatingBulkRequest<BulkOperationWrapper,
         logFailure.accept(failures.build(), failure);
     }
 
+    private void executeConsumer(final BulkOperationWrapper bulkOperationWrapper, final BulkResponseItem bulkResponseItem) {
+        if (bulkOperationWrapper.getBulkResponseItemConsumer() == null) {
+            return;
+        }
+
+        bulkOperationWrapper.getBulkResponseItemConsumer().accept(bulkResponseItem);
+    }
+
 }
diff --git a/...main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java b/...main/java/org/opensearch/dataprepper/plugins/sink/opensearch/ConnectionConfiguration.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.dataprepper.plugins.sink.opensearch;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.http.HttpHost;
 import org.apache.http.HttpRequestInterceptor;
@@ -356,6 +357,7 @@ private void attachUserCredentials(final RestClientBuilder restClientBuilder) {
     );
   }
 
+  @JsonIgnore
   private void setHttpProxyIfApplicable(final HttpAsyncClientBuilder httpClientBuilder) {
     proxy.ifPresent(
             p -> {
@@ -475,6 +477,7 @@ private static TrustManager[] createTrustManagers(final Path certPath) {
     }
   }
 
+  @JsonIgnore
   private void setHttpProxyIfApplicable(final ApacheHttpClient.Builder apacheHttpClientBuilder) {
     proxy.ifPresent(
             p -> {