Updated sigV4 signing library for gremlin and downgrade gremlin drive…

…r to Neptune supported version
awslabs · Apr 22, 2024 · 055280b · 055280b
1 parent 8e259f3
commit 055280b
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 6 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.2</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
@@ -21,9 +21,11 @@
 
 import com.amazonaws.athena.connectors.neptune.propertygraph.NeptuneGremlinConnection;
 import com.amazonaws.athena.connectors.neptune.rdf.NeptuneSparqlConnection;
+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;
@@ -45,7 +47,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();
@@ -77,7 +91,7 @@ public static NeptuneConnection createConnection(java.util.Map<String, String> c
                 throw new IllegalArgumentException("Unsupported graphType: " + graphType);
         }
     }
-    
+
     public String getNeptuneEndpoint()
     {
         return this.neptuneEndpoint;

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();