Add semantic enrichment support for OpenSearch sink index creation (opensearch-project#6771)

* Add semantic enrichment support for OpenSearch sink index creation

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

* Fixed license headers

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

* Addressed review comments. Modified to use AWS SDK instead of http endpoint for creating semantic enrichment enabled indexes

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

* Addressed review comments

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

* Fixed failing end2end tests by creating a SemanticFieldMapping as a separate class

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

* Fixed failing license checks. Addressed comments

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

* Fixed failing license checks. Addressed comments

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

---------

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

Author: kkondaka

build.gradle

@@ -285,7 +285,7 @@ subprojects {
 
 configure(subprojects.findAll {it.name != 'data-prepper-api'}) {
     dependencies {
-        implementation platform('software.amazon.awssdk:bom:2.30.23')
+        implementation platform('software.amazon.awssdk:bom:2.39.0')
         implementation 'jakarta.validation:jakarta.validation-api:3.0.2'
     }
 }
@@ -385,4 +385,4 @@ coreProjects.each { coreProject ->
     def releaseTask = tasks.named('release').get()
     releaseTask.dependsOn assembleTasks
     releaseTask.dependsOn publishTasks
-}
+}

data-prepper-plugins/opensearch/build.gradle

@@ -35,6 +35,7 @@ dependencies {
     implementation 'io.micrometer:micrometer-core'
     implementation 'software.amazon.awssdk:s3'
     implementation 'software.amazon.awssdk:opensearchserverless'
+    implementation 'software.amazon.awssdk:opensearch'
     implementation libs.commons.lang3
     implementation libs.caffeine
     implementation 'software.amazon.awssdk:apache-client'

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

@@ -1,6 +1,10 @@
 /*
  * 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;
@@ -120,7 +124,7 @@ public class ConnectionConfiguration {
   private final byte[] clientCertContent;
   private final byte[] clientKeyContent;
 
-  List<String> getHosts() {
+  public List<String> getHosts() {
     return hosts;
   }
 
@@ -131,12 +135,12 @@ String getUsername() {
   String getApitoken() {
     return apitoken;
   }
-
+  
   String getPassword() {
     return password;
   }
 
-  boolean isAwsSigv4() {
+  public boolean isAwsSigv4() {
     return awsSigv4;
   }
 

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

@@ -35,6 +35,7 @@
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexManagerFactory;
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexTemplateAPIWrapper;
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexTemplateAPIWrapperFactory;
+import org.opensearch.dataprepper.plugins.sink.opensearch.index.SemanticEnrichmentIndexManager;
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.IndexType;
 import org.opensearch.dataprepper.plugins.sink.opensearch.index.TemplateStrategy;
 import org.opensearch.dataprepper.plugins.source.opensearch.configuration.ServerlessOptions;
@@ -156,6 +157,13 @@ private void doInitializeInternal() throws IOException {
 
     maybeUpdateServerlessNetworkPolicy();
 
+    // Create index with semantic enrichment via AWS control plane (AOSS or managed domain) if configured.
+    new SemanticEnrichmentIndexManager(awsCredentialsSupplier).maybeCreateIndex(
+            connectionConfiguration,
+            openSearchSinkConfig.getIndexConfiguration().getSemanticEnrichmentConfig(),
+            openSearchSinkConfig.getIndexConfiguration().getSemanticEnrichmentResourceName(),
+            configuredIndexAlias);
+
     indexManager.setupIndex();
 
     ingester.initialize();

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

@@ -1,12 +1,17 @@
 /*
  * 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.Size;
+import org.opensearch.dataprepper.plugins.sink.opensearch.index.SemanticEnrichmentConfig;
 import software.amazon.awssdk.regions.Region;
 
 import java.util.Map;
@@ -34,9 +39,15 @@ public class AwsAuthenticationConfiguration {
     @JsonProperty("serverless")
     private boolean serverless = false;
 
+    @JsonProperty("resource_name")
+    private String resourceName;
+
     @JsonProperty("serverless_options")
     private ServerlessOptions serverlessOptions;
 
+    @JsonProperty("semantic_enrichment")
+    private SemanticEnrichmentConfig semanticEnrichmentConfig;
+
     public String getAwsStsRoleArn() {
         return awsStsRoleArn;
     }
@@ -57,8 +68,16 @@ public boolean isServerlessCollection() {
         return serverless;
     }
 
+    public String getResourceName() {
+        return resourceName;
+    }
+
     public ServerlessOptions getServerlessOptions() {
         return serverlessOptions;
     }
+
+    public SemanticEnrichmentConfig getSemanticEnrichmentConfig() {
+        return semanticEnrichmentConfig;
+    }
 }
 

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

@@ -1,6 +1,10 @@
 /*
  * 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.index;
@@ -103,8 +107,9 @@ public class IndexConfiguration {
     private final String versionExpression;
     private final VersionType versionType;
     private final boolean normalizeIndex;
-
     private final ScriptConfiguration scriptConfiguration;
+    private final SemanticEnrichmentConfig semanticEnrichmentConfig;
+    private final String semanticEnrichmentResourceName;
 
     private final String queryWhen;
 
@@ -134,6 +139,8 @@ private IndexConfiguration(final Builder builder) {
         this.versionType = builder.versionType;
         this.normalizeIndex = builder.normalizeIndex;
         this.queryOnBulkFailures = builder.queryOnIndexingFailure;
+        this.semanticEnrichmentConfig = builder.semanticEnrichmentConfig;
+        this.semanticEnrichmentResourceName = builder.semanticEnrichmentResourceName;
 
         determineTemplateType(builder);
 
@@ -179,6 +186,7 @@ private IndexConfiguration(final Builder builder) {
         this.actions = builder.actions;
         this.scriptConfiguration = builder.scriptConfiguration;
         this.documentRootKey = builder.documentRootKey;
+
         this.queryWhen = builder.queryWhen;
         this.queryTerm = builder.queryTerm;
         this.queryActionOnFound = builder.actionOnFound;
@@ -282,6 +290,12 @@ public static IndexConfiguration readIndexConfig(final OpenSearchSinkConfig open
                 .withDocumentRootKey(openSearchSinkConfig.getDocumentRootKey())
                 .withDistributionVersion(openSearchSinkConfig.getDistributionVersion());
 
+        final AwsAuthenticationConfiguration awsConfig = openSearchSinkConfig.getAwsAuthenticationOptions();
+        if (awsConfig != null && awsConfig.getSemanticEnrichmentConfig() != null) {
+            builder = builder.withSemanticEnrichmentConfig(awsConfig.getSemanticEnrichmentConfig());
+            builder = builder.withSemanticEnrichmentResourceName(awsConfig.getResourceName());
+        }
+
 
         final String versionExpression = openSearchSinkConfig.getVersionExpression();
         builder = builder.withVersionExpression(versionExpression);
@@ -329,6 +343,13 @@ public static IndexConfiguration readIndexConfig(final OpenSearchSinkConfig open
         return builder.build();
     }
 
+    public SemanticEnrichmentConfig getSemanticEnrichmentConfig() {
+        return semanticEnrichmentConfig;
+    }
+
+    public String getSemanticEnrichmentResourceName() {
+        return semanticEnrichmentResourceName;
+    }
 
     public IndexType getIndexType() {
         return indexType;
@@ -541,7 +562,9 @@ public static class Builder {
         private Duration queryDuration;
         private boolean queryOnIndexingFailure;
 
+        private SemanticEnrichmentConfig semanticEnrichmentConfig;
         private Integer queryAsyncLimit;
+        private String semanticEnrichmentResourceName;
 
         public Builder withIndexAlias(final String indexAlias) {
             checkArgument(indexAlias != null, "indexAlias cannot be null.");
@@ -782,6 +805,16 @@ public Builder withQueryAsyncLimit(final Integer queryAsyncLimit) {
             return this;
         }
 
+        public Builder withSemanticEnrichmentConfig(final SemanticEnrichmentConfig semanticEnrichmentConfig) {
+            this.semanticEnrichmentConfig = semanticEnrichmentConfig;
+            return this;
+        }
+
+        public Builder withSemanticEnrichmentResourceName(final String resourceName) {
+            this.semanticEnrichmentResourceName = resourceName;
+            return this;
+        }
+
         public IndexConfiguration build() {
             return new IndexConfiguration(this);
         }

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

@@ -0,0 +1,48 @@
+/*
+ * 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.index;
+
+import java.net.URI;
+import java.util.List;
+
+public class OpenSearchEndpointIdentifier {
+
+    private OpenSearchEndpointIdentifier() {
+    }
+
+    public static String extractCollectionId(final List<String> hosts) {
+        final String hostname = getHostname(hosts);
+        return hostname.split("\\.")[0];
+    }
+
+    public static String extractDomainName(final List<String> hosts) {
+        final String hostname = getHostname(hosts);
+        final String prefix = hostname.split("\\.")[0];
+        final String withoutSearchPrefix = prefix.replaceFirst("^(search-|vpc-)", "");
+        final int lastHyphen = withoutSearchPrefix.lastIndexOf('-');
+        if (lastHyphen <= 0) {
+            throw new IllegalArgumentException(
+                    "Unable to extract domain name from host: " + hostname +
+                            ". Please set the 'domain_name' option in semantic_enrichment config.");
+        }
+        return withoutSearchPrefix.substring(0, lastHyphen);
+    }
+
+    static String getHostname(final List<String> hosts) {
+        if (hosts == null || hosts.isEmpty()) {
+            throw new IllegalArgumentException("Hosts list is empty, cannot extract endpoint identifier");
+        }
+        final String hostname = URI.create(hosts.get(0)).getHost();
+        if (hostname == null) {
+            throw new IllegalArgumentException("Unable to parse hostname from: " + hosts.get(0));
+        }
+        return hostname;
+    }
+}

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

@@ -0,0 +1,24 @@
+/*
+ * 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.index;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.List;
+
+public class SemanticEnrichmentConfig {
+
+    @JsonProperty("fields")
+    private List<SemanticFieldMapping> fields;
+
+    public List<SemanticFieldMapping> getFields() {
+        return fields;
+    }
+}