Skip to content

Commit

Permalink
logging update and clean up
Browse files Browse the repository at this point in the history
  • Loading branch information
@xiazcy @AbdulR3hman
xiazcy authored and AbdulR3hman committed Mar 1, 2024
1 parent 6e821ae commit 51400a7
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 10 deletions.
9 changes: 8 additions & 1 deletion athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneConnection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
*/
package com.amazonaws.athena.connectors.neptune;

import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
Expand All @@ -27,10 +28,13 @@
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class NeptuneConnection
{
private static Cluster cluster = null;
private static final Logger logger = LoggerFactory.getLogger(NeptuneConnection.class);

private String neptuneEndpoint;
private String neptunePort;
Expand All @@ -45,14 +49,17 @@ protected NeptuneConnection(String neptuneEndpoint, String neptunePort, boolean
.enableSsl(true);

if (enabledIAM) {
logger.info("Connecting with IAM auth to https://" + neptuneEndpoint + ":" + neptunePort + " in " + region);
final AWSCredentialsProvider awsCredentialsProvider = new DefaultAWSCredentialsProviderChain();
builder.handshakeInterceptor(r ->
{
try {
NeptuneNettyHttpSigV4Signer sigV4Signer =
new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
new NeptuneNettyHttpSigV4Signer(region, awsCredentialsProvider);
sigV4Signer.signRequest(r);
}
catch (NeptuneSigV4SignerException e) {
logger.error("SIGV4 exception", e);
throw new RuntimeException("Exception occurred while signing the request", e);
}
return r;
Expand Down
9 changes: 8 additions & 1 deletion .../java/com/amazonaws/athena/connectors/neptune/propertygraph/NeptuneGremlinConnection.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
package com.amazonaws.athena.connectors.neptune.propertygraph;

import com.amazonaws.athena.connectors.neptune.NeptuneConnection;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
import com.amazonaws.neptune.auth.NeptuneNettyHttpSigV4Signer;
import com.amazonaws.neptune.auth.NeptuneSigV4SignerException;
Expand All @@ -28,9 +29,12 @@
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;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class NeptuneGremlinConnection extends NeptuneConnection
{
private static final Logger logger = LoggerFactory.getLogger(NeptuneGremlinConnection.class);
private static Cluster cluster = null;

public NeptuneGremlinConnection(String neptuneEndpoint, String neptunePort, boolean enabledIAM, String region)
Expand All @@ -42,14 +46,17 @@ public NeptuneGremlinConnection(String neptuneEndpoint, String neptunePort, bool
.enableSsl(true);

if (enabledIAM) {
logger.info("Connecting with IAM auth to https://" + neptuneEndpoint + ":" + neptunePort + " in " + region);
final AWSCredentialsProvider awsCredentialsProvider = new DefaultAWSCredentialsProviderChain();
builder.handshakeInterceptor(r ->
{
try {
NeptuneNettyHttpSigV4Signer sigV4Signer =
new NeptuneNettyHttpSigV4Signer(region, new DefaultAWSCredentialsProviderChain());
new NeptuneNettyHttpSigV4Signer(region, awsCredentialsProvider);
sigV4Signer.signRequest(r);
}
catch (NeptuneSigV4SignerException e) {
logger.error("SIGV4 exception", e);
throw new RuntimeException("Exception occurred while signing the request", e);
}
return r;
Expand Down
16 changes: 8 additions & 8 deletions ...m/amazonaws/athena/connectors/neptune/propertygraph/rowwriters/CustomSchemaRowWriter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ public static void writeRowTemplate(RowWriterBuilder rowWriterBuilder, Field fie

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate BIT*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");

if (fieldValue.getClass().equals(Boolean.class)) {
Boolean booleanValue = Boolean.parseBoolean(fieldValue.toString());
Expand Down Expand Up @@ -114,7 +114,7 @@ public static void writeRowTemplate(RowWriterBuilder rowWriterBuilder, Field fie
else {
Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate VARCHAR*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");

if (fieldValue != null) {
if (fieldValue.getClass().equals(String.class)) {
Expand Down Expand Up @@ -146,7 +146,7 @@ else if (fieldValue instanceof ArrayList) {

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate DATEMILLI*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");
if (fieldValue.getClass().equals(Date.class)) {
value.value = ((Date) fieldValue).getTime();
value.isSet = 1;
Expand All @@ -170,7 +170,7 @@ else if (fieldValue instanceof ArrayList) {

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate INT*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");
if (fieldValue.getClass().equals(Integer.class)) {
value.value = Integer.parseInt(fieldValue.toString());
value.isSet = 1;
Expand All @@ -194,7 +194,7 @@ else if (fieldValue instanceof ArrayList) {

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate BIGINT*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");
if (fieldValue.getClass().equals(Long.class)) {
value.value = Long.parseLong(fieldValue.toString());
value.isSet = 1;
Expand All @@ -218,7 +218,7 @@ else if (fieldValue instanceof ArrayList) {

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate FLOAT4*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");
if (fieldValue.getClass().equals(Float.class)) {
value.value = Float.parseFloat(fieldValue.toString());
value.isSet = 1;
Expand All @@ -242,7 +242,7 @@ else if (fieldValue instanceof ArrayList) {

Object fieldValue = obj.get(fieldName);
logger.debug("writeRowTemplate FLOAT8*" + field.getName() + "*" + minorType + "*"
+ fieldValue + "*" + (fieldValue == null ? "" : fieldValue.getClass()) + "*");
+ (fieldValue == null ? "" : fieldValue.getClass()) + "*");
if (fieldValue.getClass().equals(Double.class)) {
value.value = Double.parseDouble(fieldValue.toString());
value.isSet = 1;
Expand All @@ -261,7 +261,7 @@ else if (fieldValue instanceof ArrayList) {
}
catch (Throwable e) {
logger.error("writeRowTemplate exception for *" + field.getName() + "*" + minorType + "*", e);
//throw new RuntimeException(e);
throw new RuntimeException(e);
}
}

Expand Down

0 comments on commit 51400a7

Please sign in to comment.