Enhance exception handling and logging error message for Neptune QPT (#…

…2018)
awslabs · Jun 13, 2024 · 3682266 · 3682266
1 parent 79a1c6b
commit 3682266
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/Constants.java b/athena-neptune/src/main/java/com/amazonaws/athena/connectors/neptune/Constants.java
@@ -50,4 +50,6 @@ protected Constants()
 
    public static final String PREFIX_KEY = "prefix_";
    public static final int PREFIX_LEN = PREFIX_KEY.length();
+
+   public static final String GREMLIN_QUERY_SUPPORT_TYPE = "valueMap";
 }
diff --git a/...neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandler.java b/...neptune/src/main/java/com/amazonaws/athena/connectors/neptune/NeptuneMetadataHandler.java
@@ -64,6 +64,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 
 import static java.util.Objects.requireNonNull;
@@ -298,12 +299,23 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
                 logger.info("NeptuneMetadataHandler doGetQueryPassthroughSchema gremlinQuery with limit: " + gremlinQuery);
                 Object object = new PropertyGraphHandler(neptuneConnection).getResponseFromGremlinQuery(graphTraversalSource, gremlinQuery);
                 GraphTraversal graphTraversalForSchema = (GraphTraversal) object;
-                Map graphTraversalObj = null;
                 if (graphTraversalForSchema.hasNext()) {
-                    graphTraversalObj = (Map) graphTraversalForSchema.next();
+                    Object responseObj = graphTraversalForSchema.next();
+                    if (responseObj instanceof Map && gremlinQuery.contains(Constants.GREMLIN_QUERY_SUPPORT_TYPE)) {
+                        logger.info("NeptuneMetadataHandler doGetQueryPassthroughSchema gremlinQuery with valueMap");
+                        Map graphTraversalObj = (Map) responseObj;
+                        schema = getSchemaFromResults(getTableResponse, graphTraversalObj, fields);
+                        return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
+                    }
+                    else {
+                        throw new RuntimeException("Unsupported gremlin query format: We are currently supporting only valueMap gremlin queries. " +
+                                "Please make sure you are using valueMap gremlin query. " +
+                                "Example for valueMap query is g.V().hasLabel(\\\"airport\\\").valueMap().limit(5)");
+                    }
+                }
+                else {
+                    throw new NoSuchElementException("No data available for gremlin query: " + gremlinQuery);
                 }
-                schema = getSchemaFromResults(getTableResponse, graphTraversalObj, fields);
-                return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
 
             case RDF:
                 String sparqlQuery = qptArguments.get(NeptuneQueryPassthrough.QUERY);
@@ -313,12 +325,14 @@ public GetTableResponse doGetQueryPassthroughSchema(BlockAllocator allocator, Ge
                 neptuneSparqlConnection.runQuery(sparqlQuery);
                 String strim = getTableResponse.getSchema().getCustomMetadata().get(Constants.SCHEMA_STRIP_URI);
                 boolean trimURI = strim == null ? false : Boolean.parseBoolean(strim);
-                Map<String, Object> resultsMap = null;
                 if (neptuneSparqlConnection.hasNext()) {
-                    resultsMap = neptuneSparqlConnection.next(trimURI);
+                    Map<String, Object> resultsMap = neptuneSparqlConnection.next(trimURI);
+                    schema = getSchemaFromResults(getTableResponse, resultsMap, fields);
+                    return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
+                }
+                else {
+                    throw new NoSuchElementException("No data available for sparql query: " + sparqlQuery);
                 }
-                schema = getSchemaFromResults(getTableResponse, resultsMap, fields);
-                return new GetTableResponse(request.getCatalogName(), tableNameObj, schema);
 
             default:
                 throw new IllegalArgumentException("Unsupported graphType: " + graphType);