Gh-3292: REST API Gremlin Timeout and Execute Endpoints

library/tinkerpop/src/test/java/uk/gov/gchq/gaffer/tinkerpop/util/modern/Person.java

@@ -17,9 +17,12 @@
 package uk.gov.gchq.gaffer.tinkerpop.util.modern;
 
 import org.apache.tinkerpop.gremlin.structure.T;
+import org.apache.tinkerpop.gremlin.structure.Vertex;
 
 import uk.gov.gchq.gaffer.commonutil.pair.Pair;
 import uk.gov.gchq.gaffer.data.element.Entity;
+import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraph;
+import uk.gov.gchq.gaffer.tinkerpop.generator.GafferPopVertexGenerator;
 
 import static uk.gov.gchq.gaffer.tinkerpop.util.modern.GafferPopModernTestUtils.AGE;
 import static uk.gov.gchq.gaffer.tinkerpop.util.modern.GafferPopModernTestUtils.NAME;
@@ -163,4 +166,8 @@ public Entity toEntity() {
                 .property(AGE, age)
                 .build();
     }
+
+    public Vertex toVertex(GafferPopGraph graph) {
+        return new GafferPopVertexGenerator(graph)._apply(toEntity());
+    }
 }

rest-api/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/config/GremlinConfig.java

@@ -16,11 +16,11 @@
 
 package uk.gov.gchq.gaffer.rest.config;
 
+import org.apache.commons.configuration2.PropertiesConfiguration;
 import org.apache.commons.configuration2.builder.fluent.Configurations;
 import org.apache.commons.configuration2.ex.ConfigurationException;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
 import org.apache.tinkerpop.gremlin.structure.Graph;
-import org.apache.tinkerpop.gremlin.structure.util.empty.EmptyGraph;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.context.annotation.Bean;
@@ -31,28 +31,54 @@
 
 @Configuration
 public class GremlinConfig {
-
     private static final Logger LOGGER = LoggerFactory.getLogger(GremlinConfig.class);
 
     /**
      * Default path to look for a GafferPop properties file if not defined in the store properties.
      */
     private static final String DEFAULT_PROPERTIES = "/gaffer/gafferpop.properties";
 
+    /**
+     * Default timeout for executing gremlin queries (2 min).
+     */
+    private static final Long DEFAULT_REQUEST_TIMEOUT = 120000L;
+
+    /**
+     * Key for GafferPop properties file to specify the timeout on gremlin queries to the REST API.
+     */
+    private static final String REQUEST_TIMEOUT_KEY = "gaffer.rest.timeout";
+
     @Bean
-    public GraphTraversalSource graphTraversalSource(final GraphFactory graphFactory) throws Exception {
-        // Determine where to look for the GafferPop properties
-        String gafferPopProperties = graphFactory.getGraph().getStoreProperties().get(GafferPopGraph.GAFFERPOP_PROPERTIES);
-        if (gafferPopProperties == null) {
-            LOGGER.warn("GafferPop properties file was not specified. Using default location: {}", DEFAULT_PROPERTIES);
-            gafferPopProperties = DEFAULT_PROPERTIES;
-        }
+    public GraphTraversalSource graphTraversalSource(final GraphFactory graphFactory) {
         // Obtain the graph traversal
-        try (Graph graph = GafferPopGraph.open(new Configurations().properties(gafferPopProperties), graphFactory.getGraph())) {
-            return graph.traversal();
+        Graph graph = GafferPopGraph.open(findPropertiesFile(graphFactory), graphFactory.getGraph());
+        return graph.traversal();
+    }
+
+    @Bean
+    public Long requestTimeout(final GraphFactory graphFactory) {
+        return findPropertiesFile(graphFactory).getLong(REQUEST_TIMEOUT_KEY, DEFAULT_REQUEST_TIMEOUT);
+    }
+
+    /**
+     * Finds and loads the correct config file for gafferpop.
+     *
+     * @param graphFactory The graph factory.
+     * @return Loaded properties from file.
+     * @throws ConfigurationException If problem loading.
+     */
+    private PropertiesConfiguration findPropertiesFile(final GraphFactory graphFactory) {
+        try {
+            // Determine where to look for the GafferPop properties
+            String gafferPopProperties = graphFactory.getGraph().getStoreProperties().get(GafferPopGraph.GAFFERPOP_PROPERTIES);
+            if (gafferPopProperties == null) {
+                LOGGER.warn("GafferPop properties file was not specified. Using default location: {}", DEFAULT_PROPERTIES);
+                gafferPopProperties = DEFAULT_PROPERTIES;
+            }
+            return new Configurations().properties(gafferPopProperties);
         } catch (final ConfigurationException e) {
-            LOGGER.error("Error loading GafferPop config, Gremlin will be unavailable: {}", e.getMessage());
-            return EmptyGraph.instance().traversal();
+            LOGGER.warn("Using default values for GafferPop, failed to load a GafferPop config: {}", e.getMessage());
+            return new PropertiesConfiguration();
         }
     }
 }

rest-api/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/config/GremlinWebSocketConfig.java

@@ -32,16 +32,18 @@ public class GremlinWebSocketConfig implements WebSocketConfigurer {
 
     private final GraphTraversalSource g;
     private final AbstractUserFactory userFactory;
+    private final Long requestTimeout;
 
     @Autowired
-    public GremlinWebSocketConfig(final GraphTraversalSource g, final AbstractUserFactory userFactory) {
+    public GremlinWebSocketConfig(final GraphTraversalSource g, final AbstractUserFactory userFactory, final Long requestTimeout) {
         this.g = g;
         this.userFactory = userFactory;
+        this.requestTimeout = requestTimeout;
     }
 
     @Override
     public void registerWebSocketHandlers(final WebSocketHandlerRegistry registry) {
-        registry.addHandler(new GremlinWebSocketHandler(g, userFactory), "/gremlin");
+        registry.addHandler(new GremlinWebSocketHandler(g, userFactory, requestTimeout), "/gremlin");
     }
 
 }

rest-api/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/controller/GremlinController.java

@@ -16,16 +16,26 @@
 
 package uk.gov.gchq.gaffer.rest.controller;
 
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import org.apache.tinkerpop.gremlin.groovy.engine.GremlinExecutor;
 import org.apache.tinkerpop.gremlin.jsr223.ConcurrentBindings;
 import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
 import org.apache.tinkerpop.gremlin.structure.Graph;
+import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONMapper;
+import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONVersion;
+import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONWriter;
+import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONXModuleV3;
 import org.apache.tinkerpop.gremlin.structure.util.empty.EmptyGraph;
 import org.json.JSONObject;
 import org.opencypher.gremlin.server.jsr223.CypherPlugin;
 import org.opencypher.gremlin.translation.CypherAst;
 import org.opencypher.gremlin.translation.translator.Translator;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpHeaders;
 import org.springframework.web.bind.annotation.PostMapping;
@@ -34,6 +44,7 @@
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
+import uk.gov.gchq.gaffer.commonutil.otel.OtelUtil;
 import uk.gov.gchq.gaffer.core.exception.GafferRuntimeException;
 import uk.gov.gchq.gaffer.exception.SerialisationException;
 import uk.gov.gchq.gaffer.jsonserialisation.JSONSerialiser;
@@ -42,11 +53,18 @@
 import uk.gov.gchq.gaffer.rest.factory.spring.AbstractUserFactory;
 import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraph;
 import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraphVariables;
+import uk.gov.gchq.koryphe.tuple.n.Tuple2;
 
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 
 import static org.springframework.http.MediaType.APPLICATION_JSON_VALUE;
 import static org.springframework.http.MediaType.TEXT_PLAIN_VALUE;
@@ -55,22 +73,41 @@
 @Tag(name = "gremlin")
 @RequestMapping("/rest/gremlin")
 public class GremlinController {
+    /**
+     * The mapper for converting to GraphSONv3.
+     */
+    public static final GraphSONMapper GRAPHSON_V3_MAPPER = GraphSONMapper.build()
+            .version(GraphSONVersion.V3_0)
+            .addCustomModule(GraphSONXModuleV3.build()).create();
+
+    /**
+     * Writer for writing GraphSONv3 output to output streams.
+     */
+    public static final GraphSONWriter GRAPHSON_V3_WRITER = GraphSONWriter.build()
+            .mapper(GRAPHSON_V3_MAPPER)
+            .wrapAdjacencyList(true).create();
 
-    // Keys for response JSON
+    // Keys for response explain JSON
     public static final String EXPLAIN_OVERVIEW_KEY = "overview";
     public static final String EXPLAIN_OP_CHAIN_KEY = "chain";
     public static final String EXPLAIN_GREMLIN_KEY = "gremlin";
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(GremlinController.class);
+    private static final String GENERAL_ERROR_MSG = "Failed to evaluate Gremlin query: ";
+
     private final ConcurrentBindings bindings = new ConcurrentBindings();
+    private final ExecutorService executorService = Context.taskWrapping(Executors.newFixedThreadPool(4));
+    private final Long requestTimeout;
     private final AbstractUserFactory userFactory;
     private final Graph graph;
     private final Map<String, Map<String, Object>> plugins = new HashMap<>();
 
     @Autowired
-    public GremlinController(final GraphTraversalSource g, final AbstractUserFactory userFactory) {
+    public GremlinController(final GraphTraversalSource g, final AbstractUserFactory userFactory, final Long requestTimeout) {
         bindings.putIfAbsent("g", g);
         graph = g.getGraph();
         this.userFactory = userFactory;
+        this.requestTimeout = requestTimeout;
         // Add cypher plugin so cypher functions can be used in queries
         plugins.put(CypherPlugin.class.getName(), new HashMap<>());
     }
@@ -87,7 +124,31 @@
         summary = "Explain a Gremlin Query",
         description = "Runs a Gremlin query and outputs an explanation of what Gaffer operations were executed on the graph")
     public String explain(@RequestHeader final HttpHeaders httpHeaders, @RequestBody final String gremlinQuery) {
-        return runGremlinAndGetExplain(gremlinQuery, httpHeaders).toString();
+        preExecuteSetUp(httpHeaders);
+        return runGremlinQuery(gremlinQuery).get1().toString();
+    }
+
+
+    /**
+     * Endpoint for running a gremlin groovy query, will respond with an output
+     * stream of GraphSONv3 JSON.
+     *
+     * @param httpHeaders  The request headers.
+     * @param gremlinQuery The gremlin groovy query.
+     * @param response The response output stream.
+     *
+     * @throws IOException If issue writing output.
+     */
+    @PostMapping(path = "/execute", consumes = TEXT_PLAIN_VALUE, produces = APPLICATION_JSON_VALUE)
+    @io.swagger.v3.oas.annotations.Operation(
+        summary = "Run a Gremlin Query",
+        description = "Runs a Gremlin query and outputs the result as GraphSONv3 JSON")
+    public void execute(@RequestHeader final HttpHeaders httpHeaders, @RequestBody final String gremlinQuery,
+            final OutputStream response) throws IOException {
+        preExecuteSetUp(httpHeaders);
+
+        // Write to output stream for response
+        GRAPHSON_V3_WRITER.writeObject(response, runGremlinQuery(gremlinQuery).get0());
     }
 
     /**
@@ -109,11 +170,38 @@
         // Translate the cypher to gremlin, always add a .toList() otherwise Gremlin wont execute it as its lazy
         final String translation = ast.buildTranslation(Translator.builder().gremlinGroovy().enableCypherExtensions().build()) + ".toList()";
 
-        JSONObject response = runGremlinAndGetExplain(translation, httpHeaders);
+        JSONObject response = runGremlinQuery(translation).get1();
         response.put(EXPLAIN_GREMLIN_KEY, translation);
         return response.toString();
     }
 
+    /**
+     * Endpoint for running a cypher query through gremlin, will respond with an
+     * output stream of GraphSONv3 JSON.
+     *
+     * @param httpHeaders The request headers.
+     * @param cypherQuery The cypher query.
+     * @param response    The response output stream.
+     *
+     * @throws IOException If issue writing output.
+     */
+    @PostMapping(path = "/cypher/execute", consumes = TEXT_PLAIN_VALUE, produces = APPLICATION_JSON_VALUE)
+    @io.swagger.v3.oas.annotations.Operation(
+        summary = "Run a Cypher Query",
+        description = "Translates a Cypher query to Gremlin and executes it returning a GraphSONv3 JSON result." +
+                      "Note will always append a '.toList()' to the translation")
+    public void cypherExecute(@RequestHeader final HttpHeaders httpHeaders, @RequestBody final String cypherQuery,
+            final OutputStream response) throws IOException {
+        preExecuteSetUp(httpHeaders);
+        final CypherAst ast = CypherAst.parse(cypherQuery);
+        // Translate the cypher to gremlin, always add a .toList() otherwise Gremlin
+        // wont execute it as its lazy
+        final String translation = ast.buildTranslation(Translator.builder().gremlinGroovy().enableCypherExtensions().build()) + ".toList()";
+
+        // Write to output stream for response
+        GRAPHSON_V3_WRITER.writeObject(response, runGremlinQuery(translation).get0());
+    }
+
     /**
      * Gets an explanation of the last chain of operations ran on a GafferPop graph.
      * This essentially shows how a Gremlin query mapped to a Gaffer operation
@@ -154,14 +242,12 @@
     }
 
     /**
-     * Executes a given Gremlin query on the graph then formats a JSON response with
-     * the executed Gaffer operations in.
+     * Do some basic pre execute set up so the graph is ready for the gremlin
+     * request to be executed.
      *
-     * @param gremlinQuery The Gremlin groovy query.
-     * @param httpHeaders The headers for the request.
-     * @return JSON explanation.
+     * @param httpHeaders Headers for user auth
      */
-    private JSONObject runGremlinAndGetExplain(final String gremlinQuery, final HttpHeaders httpHeaders) {
+    private void preExecuteSetUp(final HttpHeaders httpHeaders) {
         // Check we actually have a graph instance to use
         GafferPopGraph gafferPopGraph;
         if (graph instanceof EmptyGraph) {
@@ -173,25 +259,70 @@
         // Hooks for user auth
         userFactory.setHttpHeaders(httpHeaders);
         graph.variables().set(GafferPopGraphVariables.USER, userFactory.createUser());
+    }
 
-        JSONObject explain = new JSONObject();
-        try (GremlinExecutor gremlinExecutor = GremlinExecutor.build()
-                .addPlugins("gremlin-groovy", plugins)
-                .globalBindings(bindings).create()) {
-            // Execute the query note this will actually run the query which we need
-            // as Gremlin will skip steps if there is no input from the previous ones
-            gremlinExecutor.eval(gremlinQuery).join();
+    /**
+     * Executes a given Gremlin query and returns the result along with an explanation.
+     *
+     * @param gremlinQuery The Gremlin groovy query.
+     * @return A pair tuple with result and explain in.
+     */
+    private Tuple2<Object, JSONObject> runGremlinQuery(final String gremlinQuery) {
+        LOGGER.info("QUERY IS: {} ", gremlinQuery);
+        GafferPopGraph gafferPopGraph = (GafferPopGraph) graph;
+
+        // OpenTelemetry hooks
+        Span span = OtelUtil.startSpan(
+            this.getClass().getName(), "Gremlin Request: " + UUID.nameUUIDFromBytes(gremlinQuery.getBytes(Charset.defaultCharset())));
+        span.setAttribute("gaffer.gremlin.query", gremlinQuery);
 
-            // Get the chain and reset the variables
-            explain = getGafferPopExplanation(gafferPopGraph);
+        // tuple to hold the result and explain
+        Tuple2<Object, JSONObject> pair = new Tuple2<>();
+        pair.put1(new JSONObject());
+
+        try (Scope scope = span.makeCurrent();
+                GremlinExecutor gremlinExecutor = getGremlinExecutor()) {
+            // Execute the query
+            Object result = gremlinExecutor.eval(gremlinQuery).join();
+
+            // Store the result and explain for returning
+            pair.put0(result);
+            pair.put1(getGafferPopExplanation(gafferPopGraph));
+
+            // Provide an debug explanation for the query that just ran
+            span.addEvent("Request complete");
+            span.setAttribute("gaffer.gremlin.explain", pair.get1().toString());
+            LOGGER.debug("{}", pair.get1());
+
+            // Reset the vars
             gafferPopGraph.setDefaultVariables((GafferPopGraphVariables) gafferPopGraph.variables());
+
         } catch (final InterruptedException e) {
             Thread.currentThread().interrupt();
+            span.setStatus(StatusCode.ERROR, e.getMessage());
+            span.recordException(e);
+            throw new GafferRuntimeException(GENERAL_ERROR_MSG + e.getMessage(), e);
         } catch (final Exception e) {
-            throw new GafferRuntimeException("Failed to evaluate Gremlin query: " + e.getMessage(), e);
+            span.setStatus(StatusCode.ERROR, e.getMessage());
+            span.recordException(e);
+            throw new GafferRuntimeException(GENERAL_ERROR_MSG + e.getMessage(), e);
         }
 
-        return explain;
+        return pair;
+    }
+
+    /**
+     * Returns a new gremlin executor. It's the responsibility of the caller to
+     * ensure it is closed.
+     *
+     * @return Gremlin executor.
+     */
+    private GremlinExecutor getGremlinExecutor() {
+        return GremlinExecutor.build()
+                .addPlugins("gremlin-groovy", plugins)
+                .evaluationTimeout(requestTimeout)
+                .executorService(executorService)
+                .globalBindings(bindings).create();
     }
 
 }