org.apache.tinkerpop.gremlin.process.traversal.Bytecode Java Examples

@Test
public void shouldSendRequestIdBytecode() {
    final UUID overrideRequestId = UUID.randomUUID();
    final Cluster cluster = TestClientFactory.build().serializer(Serializers.GRAPHSON_V3D0).create();
    final Client client = Mockito.spy(cluster.connect().alias("g"));
    Mockito.when(client.alias("g")).thenReturn(client);
    GraphTraversalSource g = AnonymousTraversalSource.traversal().withRemote(DriverRemoteConnection.using(client));
    g.with(Tokens.REQUEST_ID, overrideRequestId).V().iterate();
    cluster.close();
    ArgumentCaptor<RequestOptions> requestOptionsCaptor = ArgumentCaptor.forClass(RequestOptions.class);
    verify(client).submitAsync(Mockito.any(Bytecode.class), requestOptionsCaptor.capture());
    RequestOptions requestOptions = requestOptionsCaptor.getValue();
    assertTrue(requestOptions.getOverrideRequestId().isPresent());
    assertEquals(overrideRequestId, requestOptions.getOverrideRequestId().get());

    ArgumentCaptor<RequestMessage> requestMessageCaptor = ArgumentCaptor.forClass(RequestMessage.class);
    verify(client).submitAsync(requestMessageCaptor.capture());
    RequestMessage requestMessage = requestMessageCaptor.getValue();
    assertEquals(overrideRequestId, requestMessage.getRequestId());
}
 

@SuppressWarnings("unused")
private String translate(Bytecode bytecode) {
    // GroovyTranslator.of("g").translate(bytecode);
    List<Instruction> steps = bytecode.getStepInstructions();
    StringBuilder sb = new StringBuilder();
    sb.append("g");
    int stepsPrint = Math.min(10, steps.size());
    for (int i = 0; i < stepsPrint; i++) {
        Instruction step = steps.get(i);
        sb.append('.').append(step);
    }
    if (stepsPrint < steps.size()) {
        sb.append("..");
    }
    return sb.toString();
}
 

@Override
protected Bytecode readValue(final Buffer buffer, final GraphBinaryReader context) throws IOException {
    final Bytecode result = new Bytecode();

    final int stepsLength = buffer.readInt();
    for (int i = 0; i < stepsLength; i++) {
        result.addStep(context.readValue(buffer, String.class, false), getInstructionArguments(buffer, context));
    }

    final int sourcesLength = buffer.readInt();
    for (int i = 0; i < sourcesLength; i++) {
        result.addSource(context.readValue(buffer, String.class, false), getInstructionArguments(buffer, context));
    }

    return result;
}
 

public static void detachElements(final Bytecode bytecode) {
    for (final Bytecode.Instruction instruction : bytecode.getInstructions()) {
        final Object[] arguments = instruction.getArguments();
        for (int i = 0; i < arguments.length; i++) {
            if (arguments[i] instanceof Bytecode)
                detachElements((Bytecode) arguments[i]);
            else if(arguments[i] instanceof List) {
                final List<Object> list = new ArrayList<>();
                for(final Object object : (List)arguments[i]) {
                    list.add( DetachedFactory.detach(object, false));
                }
                arguments[i] = list;
            }
            else
                arguments[i] = DetachedFactory.detach(arguments[i], false);
        }
    }
}
 

@Override
public Bytecode.Binding deserialize(final JsonParser jsonParser, final DeserializationContext deserializationContext) throws IOException, JsonProcessingException {
    String k = null;
    Object v = null;

    while (jsonParser.nextToken() != JsonToken.END_OBJECT) {
        if (jsonParser.getCurrentName().equals(GraphSONTokens.KEY)) {
            jsonParser.nextToken();
            k = jsonParser.getText();
        } else if (jsonParser.getCurrentName().equals(GraphSONTokens.VALUE)) {
            jsonParser.nextToken();
            v = deserializationContext.readValue(jsonParser, Object.class);
        }
    }
    return new Bytecode.Binding<>(k, v);
}
 

@Test
public void shouldSendUserAgentBytecode() {
    final Cluster cluster = TestClientFactory.build().serializer(Serializers.GRAPHSON_V3D0).create();
    final Client client = Mockito.spy(cluster.connect().alias("g"));
    Mockito.when(client.alias("g")).thenReturn(client);
    GraphTraversalSource g = AnonymousTraversalSource.traversal().withRemote(DriverRemoteConnection.using(client));
    g.with(Tokens.ARGS_USER_AGENT, "test").V().iterate();
    cluster.close();
    ArgumentCaptor<RequestOptions> requestOptionsCaptor = ArgumentCaptor.forClass(RequestOptions.class);
    verify(client).submitAsync(Mockito.any(Bytecode.class), requestOptionsCaptor.capture());
    RequestOptions requestOptions = requestOptionsCaptor.getValue();
    assertEquals("test", requestOptions.getUserAgent().get());

    ArgumentCaptor<RequestMessage> requestMessageCaptor = ArgumentCaptor.forClass(RequestMessage.class);
    verify(client).submitAsync(requestMessageCaptor.capture());
    RequestMessage requestMessage = requestMessageCaptor.getValue();
    assertEquals("test", requestMessage.getArgs().getOrDefault(Tokens.ARGS_USER_AGENT, null));
}
 

protected static RequestOptions getRequestOptions(final Bytecode bytecode) {
    final Iterator<OptionsStrategy> itty = BytecodeUtil.findStrategies(bytecode, OptionsStrategy.class);
    final RequestOptions.Builder builder = RequestOptions.build();
    while (itty.hasNext()) {
        final OptionsStrategy optionsStrategy = itty.next();
        final Map<String,Object> options = optionsStrategy.getOptions();
        if (options.containsKey(ARGS_EVAL_TIMEOUT))
            builder.timeout((long) options.get(ARGS_EVAL_TIMEOUT));
        if (options.containsKey(REQUEST_ID))
            builder.overrideRequestId((UUID) options.get(REQUEST_ID));
        if (options.containsKey(ARGS_BATCH_SIZE))
            builder.batchSize((int) options.get(ARGS_BATCH_SIZE));
        if (options.containsKey(ARGS_USER_AGENT))
            builder.userAgent((String) options.get(ARGS_USER_AGENT));
    }
    return builder.create();
}
 

@Override
public CompletableFuture<ResultSet> submitAsync(final Bytecode bytecode, final RequestOptions options) {
    try {
        // need to call buildMessage() right away to get client specific configurations, that way request specific
        // ones can override as needed
        final RequestMessage.Builder request = buildMessage(RequestMessage.build(Tokens.OPS_BYTECODE)
                                                                          .processor("traversal")
                                                                          .addArg(Tokens.ARGS_GREMLIN, bytecode));

        // apply settings if they were made available
        options.getBatchSize().ifPresent(batchSize -> request.add(Tokens.ARGS_BATCH_SIZE, batchSize));
        options.getTimeout().ifPresent(timeout -> request.add(Tokens.ARGS_EVAL_TIMEOUT, timeout));
        options.getOverrideRequestId().ifPresent(request::overrideRequestId);
        options.getUserAgent().ifPresent(userAgent -> request.add(Tokens.ARGS_USER_AGENT, userAgent));

        return submitAsync(request.create());
    } catch (Exception ex) {
        throw new RuntimeException(ex);
    }
}
 

@Override
public Bytecode.Binding deserialize(final JsonParser jsonParser, final DeserializationContext deserializationContext) throws IOException, JsonProcessingException {
    String k = null;
    Object v = null;

    while (jsonParser.nextToken() != JsonToken.END_OBJECT) {
        if (jsonParser.getCurrentName().equals(GraphSONTokens.KEY)) {
            jsonParser.nextToken();
            k = jsonParser.getText();
        } else if (jsonParser.getCurrentName().equals(GraphSONTokens.VALUE)) {
            jsonParser.nextToken();
            v = deserializationContext.readValue(jsonParser, Object.class);
        }
    }
    return new Bytecode.Binding<>(k, v);
}
 

@Override
public Traversal.Admin eval(final Bytecode bytecode, final Bindings bindings, final String traversalSource) throws ScriptException {
    // these validations occur before merging in bytecode bindings which will override existing ones. need to
    // extract the named traversalsource prior to that happening so that bytecode bindings can share the same
    // namespace as global bindings (e.g. traversalsources and graphs).
    if (traversalSource.equals(HIDDEN_G))
        throw new IllegalArgumentException("The traversalSource cannot have the name " + HIDDEN_G + " - it is reserved");

    if (bindings.containsKey(HIDDEN_G))
        throw new IllegalArgumentException("Bindings cannot include " + HIDDEN_G + " - it is reserved");

    if (!bindings.containsKey(traversalSource))
        throw new IllegalArgumentException("The bindings available to the ScriptEngine do not contain a traversalSource named: " + traversalSource);

    final Object b = bindings.get(traversalSource);
    if (!(b instanceof TraversalSource))
        throw new IllegalArgumentException(traversalSource + " is of type " + b.getClass().getSimpleName() + " and is not an instance of TraversalSource");

    final Bindings inner = new SimpleBindings();
    inner.putAll(bindings);
    inner.putAll(bytecode.getBindings());
    inner.put(HIDDEN_G, b);
    org.apache.tinkerpop.gremlin.process.traversal.Script script = GroovyTranslator.of(HIDDEN_G, typeTranslator).translate(bytecode);
    script.getParameters().ifPresent(inner::putAll);
    return (Traversal.Admin) this.eval(script.getScript(), inner);
}
 

private static <O extends OutputShim> void writeInstructions(final KryoShim<?, O> kryo, final O output,
                                                             final List<Bytecode.Instruction> instructions) {
    output.writeInt(instructions.size());
    for (Bytecode.Instruction inst : instructions) {
        output.writeString(inst.getOperator());
        kryo.writeObject(output, inst.getArguments());
    }
}
 

private static <O extends OutputShim> void writeInstructions(final KryoShim<?, O> kryo, final O output,
                                                             final List<Bytecode.Instruction> instructions) {
    output.writeInt(instructions.size());
    for (Bytecode.Instruction inst : instructions) {
        output.writeString(inst.getOperator());
        kryo.writeObject(output, inst.getArguments());
    }
}
 

@Override
public <E> CompletableFuture<RemoteTraversal<?, E>> submitAsync(final Bytecode bytecode) throws RemoteConnectionException {
    try {
        return client.submitAsync(bytecode, getRequestOptions(bytecode)).thenApply(rs -> new DriverRemoteTraversal<>(rs, client, attachElements, conf));
    } catch (Exception ex) {
        throw new RemoteConnectionException(ex);
    }
}
 

@Override
public void serialize(final Bytecode.Binding binding, final JsonGenerator jsonGenerator, final SerializerProvider serializerProvider)
        throws IOException {
    jsonGenerator.writeStartObject();
    jsonGenerator.writeStringField(GraphSONTokens.KEY, binding.variable());
    jsonGenerator.writeObjectField(GraphSONTokens.VALUE, binding.value());
    jsonGenerator.writeEndObject();
}
 

@Override
public Bytecode deserialize(final JsonParser jsonParser, final DeserializationContext deserializationContext) throws IOException, JsonProcessingException {
    final Bytecode bytecode = new Bytecode();

    while (jsonParser.nextToken() != JsonToken.END_OBJECT) {
        final String current = jsonParser.getCurrentName();
        if (current.equals(GraphSONTokens.SOURCE) || current.equals(GraphSONTokens.STEP)) {
            jsonParser.nextToken();

            while (jsonParser.nextToken() != JsonToken.END_ARRAY) {

                // there should be a list now and the first item in the list is always string and is the step name
                // skip the start array
                jsonParser.nextToken();

                final String stepName = jsonParser.getText();

                // iterate through the rest of the list for arguments until it gets to the end
                final List<Object> arguments = new ArrayList<>();
                while (jsonParser.nextToken() != JsonToken.END_ARRAY) {
                    // we don't know the types here, so let the deserializer figure that business out
                    arguments.add(deserializationContext.readValue(jsonParser, Object.class));
                }

                // if it's not a "source" then it must be a "step"
                if (current.equals(GraphSONTokens.SOURCE))
                    bytecode.addSource(stepName, arguments.toArray());
                else
                    bytecode.addStep(stepName, arguments.toArray());
            }
        }
    }
    return bytecode;
}
 

@Parameterized.Parameters(name = "Value={0}")
public static Collection input() {
    final Bytecode.Binding b = new Bytecode.Binding(null, "b");

    final ReferenceVertex vertex = new ReferenceVertex("a vertex", null);

    final Bytecode bytecode = new Bytecode();
    bytecode.addStep(null);

    final BulkSet<Object> bulkSet = new BulkSet<>();
    bulkSet.add(vertex, 1L);

    final MutableMetrics metrics = new MutableMetrics("a metric", null);

    final Tree<Vertex> tree = new Tree<>();
    tree.put(vertex, null);

    // Provide instances that are malformed for serialization to fail
    return Arrays.asList(
            b,
            vertex,
            Collections.singletonMap("one", b),
            bulkSet,
            bytecode,
            Collections.singletonList(vertex),
            new ReferenceEdge("an edge", null, vertex, vertex),
            Lambda.supplier(null),
            metrics,
            new DefaultTraversalMetrics(1L, Collections.singletonList(metrics)),
            new DefaultRemoteTraverser<>(new Object(), 1L),
            tree,
            new ReferenceVertexProperty<>("a prop", null, "value"),
            new InvalidPath()
    );
}
 

/**
 * A version of {@link #submit(Bytecode)} which provides the ability to set per-request options.
 *
 * @param bytecode request in the form of gremlin {@link Bytecode}
 * @param options for the request
 *
 * @see #submit(Bytecode)
 */
public ResultSet submit(final Bytecode bytecode, final RequestOptions options) {
    try {
        return submitAsync(bytecode, options).get();
    } catch (UnsupportedOperationException uoe) {
        throw uoe;
    } catch (Exception ex) {
        throw new RuntimeException(ex);
    }
}
 

@Override
protected void writeValue(final Bytecode value, final Buffer buffer, final GraphBinaryWriter context) throws IOException {
    final List<Bytecode.Instruction> steps = value.getStepInstructions();
    final List<Bytecode.Instruction> sources = value.getSourceInstructions();
    // 2 buffers for the length + plus 2 buffers per each step and source

    writeInstructions(buffer, context, steps);
    writeInstructions(buffer, context, sources);
}
 

private void writeInstructions(final Buffer buffer, final GraphBinaryWriter context,
                               final List<Bytecode.Instruction> instructions) throws IOException {

    context.writeValue(instructions.size(), buffer, false);

    for (Bytecode.Instruction instruction : instructions) {
        context.writeValue(instruction.getOperator(), buffer, false);
        fillArgumentsBuffer(instruction.getArguments(), buffer, context);
    }
}
 

public static Optional<String> getLambdaLanguage(final Bytecode bytecode) {
    for (final Bytecode.Instruction instruction : bytecode.getInstructions()) {
        for (Object object : instruction.getArguments()) {
            if (object instanceof Lambda)
                return Optional.of(((Lambda) object).getLambdaLanguage());
            else if (object instanceof Bytecode) {
                final Optional<String> temp = BytecodeHelper.getLambdaLanguage((Bytecode) object);
                if (temp.isPresent())
                    return temp;
            }
        }
    }
    return Optional.empty();
}
 

public static void removeBindings(final Bytecode bytecode) {
    for (final Bytecode.Instruction instruction : bytecode.getInstructions()) {
        final Object[] arguments = instruction.getArguments();
        for (int i = 0; i < arguments.length; i++) {
            if (arguments[i] instanceof Bytecode.Binding)
                arguments[i] = ((Bytecode.Binding) arguments[i]).value();
            else if (arguments[i] instanceof Bytecode)
                removeBindings((Bytecode) arguments[i]);
        }
    }
}
 

/**
 * Parses {@link Bytecode} to find {@link TraversalStrategy} objects in the source instructions.
 */
public static <A extends TraversalStrategy> Iterator<A> findStrategies(final Bytecode bytecode, final Class<A> clazz) {
    return IteratorUtils.map(
            IteratorUtils.filter(bytecode.getSourceInstructions().iterator(),
                    s -> s.getOperator().equals(TraversalSource.Symbols.withStrategies) && clazz.isAssignableFrom(s.getArguments()[0].getClass())),
            os -> (A) os.getArguments()[0]);
}
 

@Override
public <E> CompletableFuture<RemoteTraversal<?, E>> submitAsync(final Bytecode bytecode) throws RemoteConnectionException {
    // default implementation for backward compatibility to 3.2.4 - this method will probably just become
    // the new submit() in 3.3.x when the deprecation is removed
    final CompletableFuture<RemoteTraversal<?, E>> promise = new CompletableFuture<>();
    try {
        promise.complete(new EmbeddedRemoteTraversal(JavaTranslator.of(g).translate(bytecode)));
    } catch (Exception t) {
        promise.completeExceptionally(t);
    }
    return promise;
}
 

@Test
public void shouldHandleBytecodeBinding() throws Exception {
    assumeThat(version, either(startsWith("v2")).or(startsWith("v3")));

    final Bytecode.Binding<String> o = new Bytecode.Binding<>("test", "testing");
    assertEquals(o, serializeDeserialize(mapper, o, Bytecode.Binding.class));
}
 

@Override
public T translate(final Bytecode bytecode) {
    try {
        final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        this.writer.writeObject(outputStream, bytecode);
        return this.wrappedTranslator.translate(this.reader.readObject(new ByteArrayInputStream(outputStream.toByteArray()), Bytecode.class));
    } catch (final Exception e) {
        throw new IllegalStateException(e.getMessage(), e);
    }
}
 

@Override
public T translate(final Bytecode bytecode) {
    try {
        final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        this.writer.writeObject(outputStream, bytecode);
        return this.wrappedTranslator.translate(this.reader.readObject(new ByteArrayInputStream(outputStream.toByteArray()), Bytecode.class));
    } catch (final Exception e) {
        throw new IllegalStateException(e.getMessage(), e);
    }
}
 

/**
 * Evaluates bytecode with bindings for a specific language into a {@link Traversal}.
 */
public Traversal.Admin eval(final Bytecode bytecode, final Bindings boundVars, final String language, final String traversalSource) throws ScriptException {
    final String lang = Optional.ofNullable(language).orElse("gremlin-groovy");

    final Bindings bindings = new SimpleBindings();
    bindings.putAll(globalBindings);
    bindings.putAll(boundVars);

    return gremlinScriptEngineManager.getEngineByName(lang).eval(bytecode, bindings, traversalSource);
}
 

@Override
public Script apply(final String traversalSource, final Object o) {
    this.script.init();
    if (o instanceof Bytecode) {
        return internalTranslate(traversalSource, (Bytecode) o);
    } else {
        return convertToScript(o);
    }
}
 

@Test
public void shouldReadWriteBinding() throws Exception {
    final String resourceName = "binding";
    assumeCompatibility(resourceName);

    final Bytecode.Binding resource = findModelEntryObject(resourceName);
    final Bytecode.Binding fromStatic = read(getCompatibility().readFromResource(resourceName), Bytecode.Binding.class);
    final Bytecode.Binding recycled = read(write(fromStatic, Bytecode.Binding.class), Bytecode.Binding.class);
    assertNotSame(fromStatic, recycled);
    assertEquals(fromStatic, recycled);
    assertEquals(resource, fromStatic);
    assertEquals(resource, recycled);
}
 

@Test
public void shouldReadWriteBytecode() throws Exception {
    final String resourceName = "bytecode";
    assumeCompatibility(resourceName);

    final Bytecode fromStatic = read(getCompatibility().readFromResource(resourceName), Bytecode.class);
    final Bytecode recycled = read(write(fromStatic, Bytecode.class), Bytecode.class);
    assertNotSame(fromStatic, recycled);
    assertEquals(fromStatic, recycled);
    // can't reasonably assert the bytecode against the original model as changes to strategies over time might
    // alter the bytecode form and then break the test. the assertions as they are ensure that the core of
    // serialization is correct by ensuring the contents of bytecode (whether they are valid for a specific version
    // or not). it seems beyond the scope of these tests to validate that bytecode itself is unchanging and fully
    // backward compatible (at least for now).
}