package com.tinkerpop.blueprints.strategy;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Function;

/**
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */

public class PartitionStrategy implements Strategy {

    private String writePartition;
    private final String partitionKey;

    public PartitionStrategy(final String partitionKey, final String partition) {
        this.writePartition = partition;
        this.partitionKey = partitionKey;
    }

    @Override
    public Function<Object[], Object[]> getPreAddVertex() {
        return (args) -> {
            final List<Object> o = new ArrayList<>(Arrays.asList(args));
            o.addAll(Arrays.asList(this.partitionKey, writePartition));
            return o.toArray();
        };
    }
}
package com.tinkerpop.blueprints.strategy;

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;

/**
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public class PartitionStrategy implements Strategy {

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

    public PartitionStrategy(final String partitionKey, final String partition) {
        this.writePartition = partition;
        this.partitionKey = partitionKey;
    }

    public String getWritePartition() {
        return writePartition;
    }

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

    public String getPartitionKey() {
        return partitionKey;
    }

    public Set<String> 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);
    }

    @Override
    public Function<Object[], Object[]> getPreAddVertex() {
        return (args) -> {
            final List<Object> o = new ArrayList<>(Arrays.asList(args));
            o.addAll(Arrays.asList(this.partitionKey, writePartition));
            return o.toArray();
        };
    }
}
Add accessors for PartitionGraph settings.