Not signinig the request for iam auth disabled cluster

xianjial · Dec 11, 2024 · 3fa15e2 · 3fa15e2
1 parent 9f21719
commit 3fa15e2
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 10 deletions.
diff --git a/...va/org/opensearch/dataprepper/plugins/source/neptune/client/NeptuneDataClientFactory.java b/...va/org/opensearch/dataprepper/plugins/source/neptune/client/NeptuneDataClientFactory.java
@@ -7,8 +7,12 @@
 import org.opensearch.dataprepper.plugins.source.neptune.configuration.AwsConfig;
 import org.opensearch.dataprepper.plugins.source.neptune.configuration.NeptuneSourceConfig;
 import software.amazon.awssdk.arns.Arn;
+import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
+import software.amazon.awssdk.core.signer.NoOpSigner;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.neptunedata.NeptunedataClient;
 import software.amazon.awssdk.services.sts.StsClient;
@@ -20,12 +24,28 @@
 
 public class NeptuneDataClientFactory {
     public static NeptunedataClient provideNeptuneDataClient(final NeptuneSourceConfig sourceConfig) {
-        final AwsConfig awsConfig = sourceConfig.getAwsConfig();
-        return NeptunedataClient
-                .builder()
-                .region(Region.of(sourceConfig.getRegion()))
-                .credentialsProvider(getAwsCredentials(Region.of(sourceConfig.getRegion()), awsConfig.getAwsStsRoleArn(), awsConfig.getAwsStsExternalId()))
-                .endpointOverride(URI.create(String.format("https://%s:%s", sourceConfig.getHost(), sourceConfig.getPort()))).build();
+        final URI endpoint = URI.create(String.format("https://%s:%s", sourceConfig.getHost(), sourceConfig.getPort()));
+        if (sourceConfig.isIamAuth()) {
+            final AwsConfig awsConfig = sourceConfig.getAwsConfig();
+            return NeptunedataClient.builder()
+                    .endpointOverride(endpoint)
+                    .region(Region.of(sourceConfig.getRegion()))
+                    .credentialsProvider(getAwsCredentials(Region.of(sourceConfig.getRegion()), awsConfig.getAwsStsRoleArn(), awsConfig.getAwsStsExternalId()))
+                    .build();
+        } else {
+            final ClientOverrideConfiguration clientOverrideConfiguration =
+                    // Do not sign the request
+                    ClientOverrideConfiguration.builder()
+                            .putAdvancedOption(SdkAdvancedClientOption.SIGNER, new NoOpSigner())
+                            .build();
+
+            return NeptunedataClient.builder()
+                    .endpointOverride(endpoint)
+                    .region(Region.of(sourceConfig.getRegion()))
+                    .overrideConfiguration(clientOverrideConfiguration)
+                    .credentialsProvider(AnonymousCredentialsProvider.create())
+                    .build();
+        }
     }
 
     private static AwsCredentialsProvider getAwsCredentials(

diff --git a/...in/java/org/opensearch/dataprepper/plugins/source/neptune/client/NeptuneStreamClient.java b/...in/java/org/opensearch/dataprepper/plugins/source/neptune/client/NeptuneStreamClient.java
@@ -38,7 +38,6 @@ public NeptuneStreamClient(final NeptuneSourceConfig config, final int batchSize
         this.retryCount = 0;
     }
 
-
     public void setStreamPosition(final long commitNum, final long opNum) {
         streamPositionInfo = new StreamPosition(commitNum, opNum);
     }

diff --git a/.../org/opensearch/dataprepper/plugins/source/neptune/configuration/NeptuneSourceConfig.java b/.../org/opensearch/dataprepper/plugins/source/neptune/configuration/NeptuneSourceConfig.java
@@ -23,7 +23,7 @@ public class NeptuneSourceConfig {
     @JsonProperty("region")
     private String region;
     @JsonProperty("iam_auth")
-    private boolean iamAuth = false;
+    private boolean iamAuth;
 
     @JsonProperty("trust_store_file_path")
     private String trustStoreFilePath;

diff --git a/...rg/opensearch/dataprepper/plugins/source/neptune/stream/StreamAcknowledgementManager.java b/...rg/opensearch/dataprepper/plugins/source/neptune/stream/StreamAcknowledgementManager.java
@@ -142,7 +142,9 @@ Optional<AcknowledgementSet> createAcknowledgementSet(final StreamCheckpoint che
     }
 
     void shutdown() {
-        monitoringTask.cancel(true);
+        if (monitoringTask != null) {
+            monitoringTask.cancel(true);
+        }
         try {
             if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
                 this.executorService.shutdownNow();

diff --git a/.../src/main/java/org/opensearch/dataprepper/plugins/source/neptune/stream/StreamWorker.java b/.../src/main/java/org/opensearch/dataprepper/plugins/source/neptune/stream/StreamWorker.java
@@ -350,7 +350,7 @@ public boolean onNeptuneStreamException(final Exception exception, final StreamP
             LOG.warn("Stream is corrupt, stopping the worker and resetting the stream.");
             this.isUnrecoverableError = true;
         } else {
-            LOG.info("Error fetching stream data, stopping processing");
+            LOG.info("Error fetching stream data, stopping processing: {}", exception.getMessage());
         }
         return false;
     }
-Original file line number
+Diff line change
@@ Expand Up @@
             this.retryCount = 0;
         }
         public void setStreamPosition(final long commitNum, final long opNum) {
             streamPositionInfo = new StreamPosition(commitNum, opNum);
         }
@@ Expand Down @@