Add generic script support for update/upsert operations in OpenSearch sink

Adds the ability to configure a script on the OpenSearch sink that gets
applied to update and upsert bulk operations. This is a generic mechanism
that passes script source and params through to OpenSearch's bulk API.
Script language is hardcoded to painless. The event document is
automatically passed as params.doc. scripted_upsert is always true
so the script runs on every write including the first create.
Script param values support ${} expression syntax for dynamic
resolution from event fields or metadata.

Configuration example:
  sink:
    - opensearch:
        hosts: ["https://localhost:9200"]
        index: "my-index"
        action: "upsert"
        document_id: "${/id}"
        script:
          source: "ctx._source.putAll(params.doc); ctx._source.source = params.table"
          params:
            table: "${getMetadata(\"table_name\")}"

Resolves #3563

Signed-off-by: Dinu John <86094133+dinujoh@users.noreply.github.com>

Author: dinujoh

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

@@ -26,6 +26,9 @@
 import org.opensearch.client.opensearch.core.bulk.IndexOperation;
 import org.opensearch.client.opensearch.core.bulk.UpdateOperation;
 import org.opensearch.client.transport.TransportOptions;
+import org.opensearch.client.json.JsonData;
+import org.opensearch.client.opensearch._types.InlineScript;
+import org.opensearch.client.opensearch._types.Script;
 import org.opensearch.common.unit.ByteSizeUnit;
 import org.opensearch.dataprepper.aws.api.AwsCredentialsSupplier;
 import org.opensearch.dataprepper.common.concurrent.BackgroundThreadFactory;
@@ -61,6 +64,7 @@
 import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.JavaClientAccumulatingUncompressedBulkRequest;
 import org.opensearch.dataprepper.plugins.sink.opensearch.bulk.SerializedJson;
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.ActionConfiguration;
+import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.ScriptConfiguration;
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.DlqConfiguration;
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.OpenSearchSinkConfig;
 import org.opensearch.dataprepper.plugins.sink.opensearch.dlq.FailedBulkOperation;
@@ -90,8 +94,10 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
@@ -115,6 +121,8 @@ public class OpenSearchSink extends AbstractSink<Record<Event>> {
   public static final String DYNAMIC_INDEX_DROPPED_EVENTS = "dynamicIndexDroppedEvents";
   public static final String INVALID_VERSION_EXPRESSION_DROPPED_EVENTS = "dynamicDocumentVersionDroppedEvents";
   private static final String PLUGIN_NAME = "opensearch";
+  private static final String PAINLESS = "painless";
+  private static final String DOC_PARAM = "doc";
 
   private static final Logger LOG = LoggerFactory.getLogger(OpenSearchSink.class);
   private static final int INITIALIZE_RETRY_WAIT_TIME_MS = 5000;
@@ -139,6 +147,7 @@ public class OpenSearchSink extends AbstractSink<Record<Event>> {
   private final String pipeline;
   private final String action;
   private final List<ActionConfiguration> actions;
+  private final ScriptConfiguration scriptConfiguration;
   private final String documentRootKey;
   private String configuredIndexAlias;
   private final ReentrantLock lock;
@@ -208,6 +217,7 @@ public OpenSearchSink(final PluginSetting pluginSetting,
     this.routing = openSearchSinkConfig.getIndexConfiguration().getRouting();
     this.action = openSearchSinkConfig.getIndexConfiguration().getAction();
     this.actions = openSearchSinkConfig.getIndexConfiguration().getActions();
+    this.scriptConfiguration = openSearchSinkConfig.getIndexConfiguration().getScriptConfiguration();
     this.documentRootKey = openSearchSinkConfig.getIndexConfiguration().getDocumentRootKey();
     this.versionType = openSearchSinkConfig.getIndexConfiguration().getVersionType();
     this.versionExpression = openSearchSinkConfig.getIndexConfiguration().getVersionExpression();
@@ -342,11 +352,12 @@ public boolean isReady() {
     return initialized;
   }
 
-  private BulkOperation getBulkOperationForAction(final String action,
+  BulkOperation getBulkOperationForAction(final String action,
                                                   final SerializedJson document,
                                                   final Long version,
                                                   final String indexName,
-                                                  final JsonNode jsonNode) {
+                                                  final JsonNode jsonNode,
+                                                  final Map<String, Object> resolvedScriptParams) {
     BulkOperation bulkOperation;
     final Optional<String> docId = document.getDocumentId();
     final Optional<String> routing = document.getRoutingField();
@@ -380,18 +391,32 @@ private BulkOperation getBulkOperationForAction(final String action,
         }
 
 
-          final UpdateOperation.Builder<Object> updateOperationBuilder = (StringUtils.equals(action.toLowerCase(), OpenSearchBulkActions.UPSERT.toString())) ?
-              new UpdateOperation.Builder<>()
+          final boolean isUpsert = StringUtils.equals(action.toLowerCase(), OpenSearchBulkActions.UPSERT.toString());
+          final UpdateOperation.Builder<Object> updateOperationBuilder = new UpdateOperation.Builder<>()
                   .index(indexName)
-                  .document(filteredJsonNode)
-                  .upsert(filteredJsonNode)
-                  .versionType(versionType)
-                  .version(version) :
-              new UpdateOperation.Builder<>()
-                  .index(indexName)
-                  .document(filteredJsonNode)
                   .versionType(versionType)
                   .version(version);
+
+          if (scriptConfiguration != null) {
+              final Map<String, JsonData> scriptParams = new HashMap<>();
+              scriptParams.put(DOC_PARAM, JsonData.of(jsonNode));
+              if (resolvedScriptParams != null) {
+                  resolvedScriptParams.forEach((k, v) -> scriptParams.put(k, JsonData.of(v)));
+              }
+              final InlineScript inlineScript = new InlineScript.Builder()
+                      .source(scriptConfiguration.getSource())
+                      .lang(PAINLESS)
+                      .params(scriptParams)
+                      .build();
+              updateOperationBuilder.script(Script.of(s -> s.inline(inlineScript)));
+              updateOperationBuilder.upsert(filteredJsonNode);
+              updateOperationBuilder.scriptedUpsert(true);
+          } else if (isUpsert) {
+              updateOperationBuilder.document(filteredJsonNode).upsert(filteredJsonNode);
+          } else {
+              updateOperationBuilder.document(filteredJsonNode);
+          }
+
           docId.ifPresent(updateOperationBuilder::id);
           routing.ifPresent(updateOperationBuilder::routing);
           bulkOperation = new BulkOperation.Builder()
@@ -535,7 +560,19 @@ public void doOutput(final Collection<Record<Event>> records) {
       BulkOperation bulkOperation;
 
       try {
-        bulkOperation = getBulkOperationForAction(eventAction, document, version, indexName, event.getJsonNode());
+        Map<String, Object> resolvedScriptParams = null;
+        if (scriptConfiguration != null && scriptConfiguration.getParams() != null) {
+          resolvedScriptParams = new HashMap<>();
+          for (final Map.Entry<String, Object> entry : scriptConfiguration.getParams().entrySet()) {
+            final Object value = entry.getValue();
+            if (value instanceof String && ((String) value).contains("${")) {
+              resolvedScriptParams.put(entry.getKey(), event.formatString((String) value, expressionEvaluator));
+            } else {
+              resolvedScriptParams.put(entry.getKey(), value);
+            }
+          }
+        }
+        bulkOperation = getBulkOperationForAction(eventAction, document, version, indexName, event.getJsonNode(), resolvedScriptParams);
       } catch (final Exception e) {
         LOG.error("An exception occurred while constructing the bulk operation for a document: ", e);
         logFailureForDlqObjects(List.of(createDlqObjectFromEvent(event, indexName, e.getMessage())), e);

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

@@ -164,6 +164,22 @@ public boolean getEnableRequestCompression() {
     @JsonProperty("actions")
     private List<ActionConfiguration> actions = null;
 
+    @Getter
+    @Valid
+    @JsonProperty("script")
+    private ScriptConfiguration scriptConfiguration = null;
+
+    @AssertTrue(message = "script can only be used with update or upsert actions.")
+    public boolean isScriptActionValid() {
+        if (scriptConfiguration == null) {
+            return true;
+        }
+        final String effectiveAction = action != null ? action : OpenSearchBulkActions.INDEX.toString();
+        return effectiveAction.contains("${") ||
+                effectiveAction.equalsIgnoreCase(OpenSearchBulkActions.UPDATE.toString()) ||
+                effectiveAction.equalsIgnoreCase(OpenSearchBulkActions.UPSERT.toString());
+    }
+
     @Getter
     @JsonProperty("document_root_key")
     private String documentRootKey = null;

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

@@ -0,0 +1,29 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.sink.opensearch.configuration;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.NotEmpty;
+import lombok.Getter;
+
+import java.util.Map;
+
+public class ScriptConfiguration {
+
+    @Getter
+    @NotEmpty
+    @JsonProperty("source")
+    private String source;
+
+    @Getter
+    @JsonProperty("params")
+    private Map<String, Object> params;
+}

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

@@ -16,6 +16,7 @@
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.OpenSearchSinkConfig;
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.ActionConfiguration;
 import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.AwsAuthenticationConfiguration;
+import org.opensearch.dataprepper.plugins.sink.opensearch.configuration.ScriptConfiguration;
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.model.QueryForExistingDocumentConfiguration;
 import org.opensearch.dataprepper.plugins.sink.opensearch.s3.FileReader;
 import org.opensearch.dataprepper.plugins.sink.opensearch.s3.S3ClientProvider;
@@ -103,6 +104,8 @@ public class IndexConfiguration {
     private final VersionType versionType;
     private final boolean normalizeIndex;
 
+    private final ScriptConfiguration scriptConfiguration;
+
     private final String queryWhen;
 
     private final Duration queryDuration;
@@ -174,6 +177,7 @@ private IndexConfiguration(final Builder builder) {
         this.ismPolicyFile = builder.ismPolicyFile;
         this.action = builder.action;
         this.actions = builder.actions;
+        this.scriptConfiguration = builder.scriptConfiguration;
         this.documentRootKey = builder.documentRootKey;
         this.queryWhen = builder.queryWhen;
         this.queryTerm = builder.queryTerm;
@@ -293,6 +297,11 @@ public static IndexConfiguration readIndexConfig(final OpenSearchSinkConfig open
             builder = builder.withAction(openSearchSinkConfig.getAction(), expressionEvaluator);
         }
 
+        final ScriptConfiguration scriptConfiguration = openSearchSinkConfig.getScriptConfiguration();
+        if (scriptConfiguration != null) {
+            builder = builder.withScriptConfiguration(scriptConfiguration);
+        }
+
         AwsAuthenticationConfiguration awsAuthenticationConfiguration = openSearchSinkConfig.getAwsAuthenticationOptions();
         if (awsAuthenticationConfiguration != null) {
             builder = builder.withServerless(awsAuthenticationConfiguration.isServerlessCollection());
@@ -372,6 +381,10 @@ public List<ActionConfiguration> getActions() {
         return actions;
     }
 
+    public ScriptConfiguration getScriptConfiguration() {
+        return scriptConfiguration;
+    }
+
     public String getS3AwsRegion() {
         return s3AwsRegion;
     }
@@ -501,6 +514,7 @@ public static class Builder {
         private Optional<String> ismPolicyFile;
         private String action;
         private List<ActionConfiguration> actions;
+        private ScriptConfiguration scriptConfiguration;
         private String s3AwsRegion;
         private String s3AwsStsRoleArn;
         private String s3AwsStsExternalId;
@@ -629,6 +643,11 @@ public Builder withActions(final List<ActionConfiguration> actions, final Expres
             return this;
         }
 
+        public Builder withScriptConfiguration(final ScriptConfiguration scriptConfiguration) {
+            this.scriptConfiguration = scriptConfiguration;
+            return this;
+        }
+
         public Builder withS3AwsRegion(final String s3AwsRegion) {
             checkNotNull(s3AwsRegion, "s3AwsRegion cannot be null");
             this.s3AwsRegion = s3AwsRegion;