com.tinkerpop.gremlin.structure.strategy.PartitionGraphStrategy Maven / Gradle / Ivy

Go to download
package com.tinkerpop.gremlin.structure.strategy;

import com.tinkerpop.gremlin.structure.Edge;
import com.tinkerpop.gremlin.structure.Element;
import com.tinkerpop.gremlin.structure.Property;
import com.tinkerpop.gremlin.structure.Vertex;
import com.tinkerpop.gremlin.util.function.STriFunction;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Function;
import java.util.function.UnaryOperator;

/**
 * A GraphStrategy which enables support for logical graph partitioning where the Graph can be blinded to different parts of the total Graph
 *
 * @author Stephen Mallette (http://stephen.genoprime.com)
 * @author Joshua Shinavier (http://fortytwo.net)
 */
public class PartitionGraphStrategy extends SubgraphStrategy {

    private String writePartition;
    private final String partitionKey;
    private final Set readPartitions = new HashSet<>();

    public PartitionGraphStrategy(final String partitionKey, final String partition) {
        super(null, null);
        this.vertexPredicate = this::testElement;
        this.edgePredicate = this::testElement;

        this.writePartition = partition;
        this.addReadPartition(partition);
        this.partitionKey = partitionKey;
    }

    private boolean testElement(final Element e) {
        final Property p = e.property(this.partitionKey);
        return p.isPresent() && this.readPartitions.contains(p.value());
    }

    public String getWritePartition() {
        return writePartition;
    }

    public void setWritePartition(final String writePartition) {
        this.writePartition = writePartition;
    }

    public String getPartitionKey() {
        return partitionKey;
    }

    public Set getReadPartitions() {
        return Collections.unmodifiableSet(readPartitions);
    }

    public void removeReadPartition(final String readPartition) {
        this.readPartitions.remove(readPartition);
    }

    public void addReadPartition(final String readPartition) {
        this.readPartitions.add(readPartition);
    }

    public void clearReadPartitions() {
        this.readPartitions.clear();
    }

    @Override
    public UnaryOperator> getAddVertexStrategy(final Strategy.Context ctx) {
        return (f) -> (keyValues) -> {
            final List