Skip to content

Commit

Permalink
test tweaks
Browse files Browse the repository at this point in the history
  • Loading branch information
@tb06904
tb06904 committed Jan 16, 2025
1 parent 1f0d16e commit 89d808f
Show file tree
Hide file tree
Showing 5 changed files with 43 additions and 41 deletions.
1 change: 0 additions & 1 deletion ...-api/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/config/GremlinWebSocketConfig.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@

package uk.gov.gchq.gaffer.rest.config;

import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
Expand Down
44 changes: 25 additions & 19 deletions rest-api/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/controller/GremlinController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,8 @@
import io.opentelemetry.context.Context;
import io.opentelemetry.context.Scope;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.apache.commons.lang3.tuple.MutablePair;
import org.apache.commons.lang3.tuple.Pair;
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;
Expand Down Expand Up @@ -55,7 +57,6 @@
import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraph;
import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraphVariables;
import uk.gov.gchq.gaffer.user.User;
import uk.gov.gchq.koryphe.tuple.n.Tuple2;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
Expand Down Expand Up @@ -125,7 +126,7 @@ public GremlinController(final GafferPopGraph graph, final AbstractUserFactory u
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 runGremlinQuery(gremlinQuery, httpHeaders).get1().toString();
return runGremlinQuery(gremlinQuery, httpHeaders).getRight().toString();
}

/**
Expand All @@ -149,7 +150,7 @@ public ResponseEntity<StreamingResponseBody> execute(
HttpStatus status = HttpStatus.OK;
StreamingResponseBody responseBody;
try {
Object result = runGremlinQuery(gremlinQuery, httpHeaders).get0();
Object result = runGremlinQuery(gremlinQuery, httpHeaders).getLeft();
// Write to output stream for response
responseBody = os -> GRAPHSON_V3_WRITER.writeObject(os, result);
} catch (final Exception e) {
Expand Down Expand Up @@ -183,7 +184,7 @@ public String cypherExplain(@RequestHeader final HttpHeaders httpHeaders, @Reque
final String translation = ast.buildTranslation(
Translator.builder().gremlinGroovy().enableCypherExtensions().build()) + ".toList()";

JSONObject response = runGremlinQuery(translation, httpHeaders).get1();
JSONObject response = runGremlinQuery(translation, httpHeaders).getRight();
response.put(EXPLAIN_GREMLIN_KEY, translation);
return response.toString();
}
Expand Down Expand Up @@ -216,7 +217,7 @@ public ResponseEntity<StreamingResponseBody> cypherExecute(
final String translation = ast.buildTranslation(
Translator.builder().gremlinGroovy().enableCypherExtensions().build()) + ".toList()";
// Run Query
Object result = runGremlinQuery(translation, httpHeaders).get0();
Object result = runGremlinQuery(translation, httpHeaders).getLeft();
// Write to output stream for response
responseBody = os -> GRAPHSON_V3_WRITER.writeObject(os, result);
} catch (final Exception e) {
Expand All @@ -238,15 +239,15 @@ public ResponseEntity<StreamingResponseBody> cypherExecute(
* query may be absent from the operation chains in the explain as it may have
* been done in the Tinkerpop framework instead.
*
* @param graph The GafferPop graph
* @param graphInstance The GafferPop graph instance.
* @return A JSON payload with an overview and full JSON representation of the
* chain in.
*/
public static JSONObject getGafferPopExplanation(final GafferPopGraph graph) {
public static JSONObject getGafferPopExplanation(final GafferPopGraph graphInstance) {
JSONObject result = new JSONObject();
// Get the last operation chain that ran
LinkedList<Operation> operations = new LinkedList<>();
((GafferPopGraphVariables) graph.variables())
((GafferPopGraphVariables) graphInstance.variables())
.getLastOperationChain()
.getOperations()
.forEach(op -> {
Expand All @@ -273,7 +274,8 @@ public static JSONObject getGafferPopExplanation(final GafferPopGraph graph) {
* Do some basic pre execute set up so the graph is ready for the gremlin
* request to be executed.
*
* @param GafferPopGraph The graph structure to use
* @param graphInstance The graph instance to bind the traversal to.
* @return The set-up {@link GremlinExecutor}.
*/
private GremlinExecutor setUpExecutor(final GafferPopGraph graphInstance) {
final ConcurrentBindings bindings = new ConcurrentBindings();
Expand All @@ -289,20 +291,23 @@ private GremlinExecutor setUpExecutor(final GafferPopGraph graphInstance) {
}

/**
* Executes a given Gremlin query and returns the result along with an explanation.
* 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.
* @param httpHeaders The headers from the request for the
* {@link AbstractUserFactory}.
* @return A pair with result left and explain right.
*/
private Tuple2<Object, JSONObject> runGremlinQuery(final String gremlinQuery, final HttpHeaders httpHeaders) {
private Pair<Object, JSONObject> runGremlinQuery(final String gremlinQuery, final HttpHeaders httpHeaders) {
// We can't reuse the existing graph instance as we need to set variables
// specific to this query only.
final GafferPopGraph graphInstance = graph.newInstance();

// Hooks for user auth
userFactory.setHttpHeaders(httpHeaders);
User user = userFactory.createUser();
graph.variables().set(GafferPopGraphVariables.USER, user);
graphInstance.variables().set(GafferPopGraphVariables.USER, user);

// OpenTelemetry hooks
Span span = OtelUtil.startSpan(
Expand All @@ -316,22 +321,23 @@ private Tuple2<Object, JSONObject> runGremlinQuery(final String gremlinQuery, fi
span.setAttribute(OtelUtil.USER_ATTRIBUTE, "unknownGremlinUser");
}

// tuple to hold the result and explain
Tuple2<Object, JSONObject> pair = new Tuple2<>();
pair.put1(new JSONObject());
// Pair to hold the result and explain
MutablePair<Object, JSONObject> pair = new MutablePair<>(null, new JSONObject());

try (Scope scope = span.makeCurrent();
GremlinExecutor gremlinExecutor = setUpExecutor(graphInstance)) {
// Execute the query
Object result = gremlinExecutor.eval(gremlinQuery).get();

// Store the result and explain for returning
pair.put0(result);
pair.put1(getGafferPopExplanation(graph));
pair.setLeft(result);
pair.setRight(getGafferPopExplanation(graphInstance));

// Provide an debug explanation for the query that just ran
span.addEvent("Request complete");
LOGGER.debug("{}", pair.get1());
if (LOGGER.isDebugEnabled()) {
LOGGER.debug("{}", pair.getRight());
}
} catch (final InterruptedException e) {
Thread.currentThread().interrupt();
span.setStatus(StatusCode.ERROR, e.getMessage());
Expand Down
16 changes: 9 additions & 7 deletions ...pi/spring-rest/src/main/java/uk/gov/gchq/gaffer/rest/handler/GremlinWebSocketHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -95,9 +95,10 @@ public class GremlinWebSocketHandler extends BinaryWebSocketHandler {
/**
* Constructor
*
* @param g The graph traversal source
* @param userFactory The user factory
* @param requestTimeout The timeout for gremlin requests
* @param graph The {@link GafferPopGraph} this is initialised to use the
* correct underlying Gaffer graph.
* @param userFactory The user factory.
* @param requestTimeout The timeout for gremlin requests.
*/
public GremlinWebSocketHandler(final GafferPopGraph graph, final AbstractUserFactory userFactory, final Long requestTimeout) {
this.graph = graph;
Expand Down Expand Up @@ -131,7 +132,8 @@ protected void handleBinaryMessage(final WebSocketSession session, final BinaryM
* Do some basic pre execute set up so the graph is ready for the gremlin
* request to be executed.
*
* @param GafferPopGraph The graph structure to use
* @param graphInstance The graph instance to bind the traversal to.
* @return The set-up {@link GremlinExecutor}.
*/
private GremlinExecutor setUpExecutor(final GafferPopGraph graphInstance) {
final ConcurrentBindings bindings = new ConcurrentBindings();
Expand Down Expand Up @@ -169,11 +171,11 @@ private ResponseMessage handleGremlinRequest(final WebSocketSession session, fin

// Execute the query
try (Scope scope = span.makeCurrent();
GremlinExecutor gremlinExecutor = setUpExecutor(graph)) {
GremlinExecutor gremlinExecutor = setUpExecutor(graphInstance)) {
// Set current headers for potential authorisation then set the user
userFactory.setHttpHeaders(session.getHandshakeHeaders());
User user = userFactory.createUser();
graph.variables().set(GafferPopGraphVariables.USER, user);
graphInstance.variables().set(GafferPopGraphVariables.USER, user);
span.setAttribute(OtelUtil.USER_ATTRIBUTE, user.getUserId());

// Run the query using the gremlin executor service
Expand All @@ -195,7 +197,7 @@ private ResponseMessage handleGremlinRequest(final WebSocketSession session, fin
// Provide an debug explanation for the query that just ran
span.addEvent("Request complete");
if (LOGGER.isDebugEnabled()) {
JSONObject gafferOperationChain = GremlinController.getGafferPopExplanation(graph);
JSONObject gafferOperationChain = GremlinController.getGafferPopExplanation(graphInstance);
LOGGER.debug("{}", gafferOperationChain);
}
// Build the response
Expand Down
11 changes: 4 additions & 7 deletions ...i/spring-rest/src/test/java/uk/gov/gchq/gaffer/rest/controller/GremlinControllerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,6 @@

package uk.gov.gchq.gaffer.rest.controller;

import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.json.JSONArray;
import org.json.JSONObject;
import org.junit.jupiter.api.Test;
Expand Down Expand Up @@ -65,9 +63,8 @@ class GremlinControllerTest {
@TestConfiguration
static class TestConfig {
@Bean
public GraphTraversalSource g() {
Graph graph = GafferPopModernTestUtils.createModernGraph(TestConfig.class, StoreType.MAP);
return graph.traversal();
public GafferPopGraph graph() {
return GafferPopModernTestUtils.createModernGraph(TestConfig.class, StoreType.MAP);
}

@Bean
Expand All @@ -85,15 +82,15 @@ public Long timeout() {
private MockMvc mockMvc;

@Autowired
private GraphTraversalSource g;
private GafferPopGraph graph;

@Test
void shouldExecuteValidGremlinQuery() throws Exception {
String gremlinString = "g.V('" + MARKO.getId() + "').toList()";

// Create the expected output
OutputStream expectedOutput = new ByteArrayOutputStream();
GremlinController.GRAPHSON_V3_WRITER.writeObject(expectedOutput, Arrays.asList(MARKO.toVertex((GafferPopGraph) g.getGraph())));
GremlinController.GRAPHSON_V3_WRITER.writeObject(expectedOutput, Arrays.asList(MARKO.toVertex((GafferPopGraph) graph)));

// When
MvcResult result = mockMvc
Expand Down
12 changes: 5 additions & 7 deletions ...ng-rest/src/test/java/uk/gov/gchq/gaffer/rest/integration/handler/GremlinWebSocketIT.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,6 @@
import org.apache.tinkerpop.gremlin.driver.Client;
import org.apache.tinkerpop.gremlin.driver.Cluster;
import org.apache.tinkerpop.gremlin.driver.Result;
import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
import org.apache.tinkerpop.gremlin.structure.Graph;
import org.apache.tinkerpop.gremlin.util.ser.GraphSONMessageSerializerV3;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
Expand All @@ -37,6 +35,7 @@

import uk.gov.gchq.gaffer.rest.factory.spring.AbstractUserFactory;
import uk.gov.gchq.gaffer.rest.factory.spring.UnknownUserFactory;
import uk.gov.gchq.gaffer.tinkerpop.GafferPopGraph;
import uk.gov.gchq.gaffer.tinkerpop.util.GafferPopTestUtil.StoreType;
import uk.gov.gchq.gaffer.tinkerpop.util.modern.GafferPopModernTestUtils;

Expand Down Expand Up @@ -70,9 +69,8 @@ static class TestConfig {

@Bean
@Profile("test")
public GraphTraversalSource g() {
Graph graph = GafferPopModernTestUtils.createModernGraph(TestConfig.class, StoreType.MAP);
return graph.traversal();
public GafferPopGraph graph() {
return GafferPopModernTestUtils.createModernGraph(TestConfig.class, StoreType.MAP);
}

@Bean
Expand All @@ -92,7 +90,7 @@ public Long timeout() {
private Integer port;

@Autowired
private GraphTraversalSource g;
private GafferPopGraph graph;

private Client client;

Expand Down Expand Up @@ -161,7 +159,7 @@ void shouldAcceptGremlinQueryUsingCustomCypherFunctions() {

// Then
assertThat(results)
.map(result -> result.getObject())
.map(Result::getObject)
.containsExactlyInAnyOrder(
String.valueOf(MARKO.getAge()),
String.valueOf(VADAS.getAge()),
Expand Down

0 comments on commit 89d808f

Please sign in to comment.