Search code examples
spring-bootgremlinamazon-neptune

Gremlin query to get nodes, edges and children in Java


For compatibility with APIs previously exposed by a Java library for Spring Boot, which I am modifying to integrate with AWS Neptune, I need to load a "node" and all its children, along with their edges (from node to children), from Neptune. However, I haven't been able to accomplish this with a single Gremlin query so far; the only way I found, which I'll outline below, involves two separate queries. Obviously, this significantly impacts performance. Is there a more elegant and optimized way to achieve the same result?

(as you can see, nodes have an attribute called entityId and a name)

@Repository
@Slf4j
public class VisibilityRepositoryGremlin {

    private final GraphTraversalSource g;

    @Autowired
    private Client client;

    @Autowired
    public VisibilityRepositoryGremlin(GraphTraversalSource g) {
        this.g = g;
    }


    public Mono<Node> findVisibleNode(UUID originEntityId, String originLabel,
                                      UUID targetEntityId, String targetLabel, boolean isPrivileged) {

        return g.V()
            .hasLabel(originLabel)
            .has(Node.ENTITY_ID_PROPERTY, originEntityId.toString())
            .repeat(outE(Node.CAN_SEE_REL_TYPE, CONTAINS_REL_TYPE)
                .has(VisibilityGroupRelationship.VISIBILITY, isPrivileged ?
                    within(VisibilityGroupRelationship.Visibility.PRIVILEGED.name(),
                           VisibilityGroupRelationship.Visibility.STANDARD.name()) :
                    within(VisibilityGroupRelationship.Visibility.STANDARD.name()))
                .otherV().dedup())
            .until(hasLabel(targetLabel)
                .has(Node.ENTITY_ID_PROPERTY, targetEntityId.toString()))
            .elementMap().fold().next()
            .stream()
            .map(VisibilityRepositoryGremlin::getNodeFromVertexProps)
            .map(vg->findById(vg.getId()))
            .findAny().orElse(Mono.error(new NotFoundException(("Entity not found"))));
    }

    

    @SuppressWarnings("unchecked")
    public Mono<Node> findById(String s) {

        List<Map<String, Object>> result= g.V().hasId(s)
            .project("visibilityGroup", "children")
            .by(elementMap())
            .by(outE().hasLabel(CONTAINS_REL_TYPE)
                .project("edge", "visibility")
                .by(inV().elementMap())
                .by("visibility")
                .fold())
            .fold().next();

        if (result.isEmpty()) return Mono.error(new NotFoundException("Not found"));

        Node vg = getNodeFromVertexProps((Map<Object, Object>)result.get(0).get("visibilityGroup"));

        List<Map<Object, Object>> childrenMaps = (List<Map<Object, Object>>)result.get(0).get("children");

        childrenMaps.forEach(map -> {
            Map<Object, Object> edgeProps = (Map<Object, Object>) map.get("edge");
            Node child = getNodeFromVertexProps(edgeProps);
            if (VisibilityGroupRelationship.Visibility.valueOf((String)map.get("visibility")) == 
                    VisibilityGroupRelationship.Visibility.PRIVILEGED)
                vg.addExplicitlyVisibleChild(child);
            else vg.addChild(child);
        });

        return Mono.just(vg);
    }

    

    private static Node getNodeFromVertexProps(Map<Object, Object> r) {
        return Node.builder()
            .id(r.get(T.id).toString())
            .entityId(UUID.fromString(r.get(Node.ENTITY_ID_PROPERTY).toString()))
            .nodeName(r.get("nodeName").toString())
            .label(r.get(T.label).toString())
            .build();
    }
}


Solution

  • You're doing a lot of extra work at the bottom of the query within the findVisibleNode method. You could combine those two into:

           g.V()
                .hasLabel(originLabel)
                .has(Node.ENTITY_ID_PROPERTY, originEntityId.toString())
                .repeat(outE(Node.CAN_SEE_REL_TYPE, CONTAINS_REL_TYPE)
                    .has(VisibilityGroupRelationship.VISIBILITY, isPrivileged ?
                        within(VisibilityGroupRelationship.Visibility.PRIVILEGED.name(),
                               VisibilityGroupRelationship.Visibility.STANDARD.name()) :
                        within(VisibilityGroupRelationship.Visibility.STANDARD.name()))
                    .otherV().dedup())
                .until(hasLabel(targetLabel)
                    .has(Node.ENTITY_ID_PROPERTY, targetEntityId.toString()))
                .project("visibilityGroup", "children")
                .by(elementMap())
                .by(outE().hasLabel(CONTAINS_REL_TYPE)
                    .project("edge", "visibility")
                    .by(inV().elementMap())
                    .by("visibility")
                    .fold())
                .fold().next();
    

    That would essentially get the same results.