Updated sigV4 signing library for gremlin and downgrade gremlin drive…

…r to Neptune supported version
awslabs · Mar 1, 2024 · 20f504d · 20f504d
1 parent 38d49be
commit 20f504d
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 5 deletions.
diff --git a/athena-neptune/pom.xml b/athena-neptune/pom.xml
@@ -9,7 +9,8 @@
     <artifactId>athena-neptune</artifactId>
     <version>2022.47.1</version>
     <properties>
-        <gremlinDriverVersion>3.7.1</gremlinDriverVersion>
+        <!-- make sure gremlin driver version stays within the Neptune supported range -->
+        <gremlinDriverVersion>3.6.5</gremlinDriverVersion>
         <neptune.sigv4.signer.version>2.4.0</neptune.sigv4.signer.version>
     </properties>
     <dependencies>

diff --git a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java
@@ -19,9 +19,11 @@
  */
 package com.amazonaws.athena.connectors.neptune;
 
+import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
+import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
+import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
 import org.apache.tinkerpop.gremlin.driver.Client;
 import org.apache.tinkerpop.gremlin.driver.Cluster;
-import org.apache.tinkerpop.gremlin.driver.SigV4WebSocketChannelizer;
 import org.apache.tinkerpop.gremlin.driver.remote.DriverRemoteConnection;
 import org.apache.tinkerpop.gremlin.process.traversal.AnonymousTraversalSource;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
@@ -43,7 +45,19 @@ protected NeptuneConnection(String neptuneEndpoint, String neptunePort, boolean
                .enableSsl(true);
 
         if (enabledIAM) {
-            builder = builder.channelizer(SigV4WebSocketChannelizer.class);
+            builder.handshakeInterceptor(r ->
+                    {
+                        try {
+                            NeptuneNettyHttpSigV4Signer sigV4Signer =
+                                    new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
+                            sigV4Signer.signRequest(r);
+                        }
+                        catch (NeptuneSigV4SignerException e) {
+                            throw new RuntimeException("Exception occurred while signing the request", e);
+                        }
+                        return r;
+                    }
+            );
         }
 
         cluster = builder.create();

diff --git a/.../java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java b/.../java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java
@@ -20,9 +20,11 @@
 package com.amazonaws.athena.connectors.neptune.propertygraph;
 
 import com.amazonaws.athena.connectors.neptune.NeptuneConnection;
+import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
+import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
+import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
 import org.apache.tinkerpop.gremlin.driver.Client;
 import org.apache.tinkerpop.gremlin.driver.Cluster;
-import org.apache.tinkerpop.gremlin.driver.SigV4WebSocketChannelizer;
 import org.apache.tinkerpop.gremlin.driver.remote.DriverRemoteConnection;
 import org.apache.tinkerpop.gremlin.process.traversal.AnonymousTraversalSource;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
@@ -40,7 +42,19 @@ public NeptuneGremlinConnection(String neptuneEndpoint, String neptunePort, bool
                .enableSsl(true);
 
         if (enabledIAM) {
-            builder = builder.channelizer(SigV4WebSocketChannelizer.class);
+            builder.handshakeInterceptor(r ->
+                    {
+                        try {
+                            NeptuneNettyHttpSigV4Signer sigV4Signer =
+                                    new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
+                            sigV4Signer.signRequest(r);
+                        }
+                        catch (NeptuneSigV4SignerException e) {
+                            throw new RuntimeException("Exception occurred while signing the request", e);
+                        }
+                        return r;
+                    }
+            );
         }
 
         cluster = builder.create();