StorageClient.java

/* Copyright (c) 2020 vesoft inc. All rights reserved.
 *
 * This source code is licensed under Apache 2.0 License.
 */
 
package com.vesoft.nebula.client.storage;
 
import com.vesoft.nebula.HostAddr;
import com.vesoft.nebula.client.graph.data.HostAddress;
import com.vesoft.nebula.client.graph.data.SSLParam;
import com.vesoft.nebula.client.meta.MetaManager;
import com.vesoft.nebula.client.storage.scan.PartScanInfo;
import com.vesoft.nebula.client.storage.scan.ScanEdgeResultIterator;
import com.vesoft.nebula.client.storage.scan.ScanVertexResultIterator;
import com.vesoft.nebula.meta.ColumnDef;
import com.vesoft.nebula.meta.Schema;
import com.vesoft.nebula.storage.EdgeProp;
import com.vesoft.nebula.storage.ScanEdgeRequest;
import com.vesoft.nebula.storage.ScanVertexRequest;
import com.vesoft.nebula.storage.VertexProp;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
 
public class StorageClient implements Serializable {
    private static final Logger LOGGER = LoggerFactory.getLogger(StorageClient.class);
 
    private final GraphStorageConnection connection;
    private StorageConnPool pool;
    private MetaManager metaManager;
    private final List<HostAddress> addresses;
    private int timeout = 10000; // ms
    private int connectionRetry = 3;
    private int executionRetry = 1;
 
    private boolean enableSSL = false;
    private SSLParam sslParam = null;
 
    private String user = null;
    private String password = null;
 
    private Map<String, String> storageAddressMapping = null;
 
    public StorageClient(String ip, int port) {
        this(Arrays.asList(new HostAddress(ip, port)));
    }
 
    public StorageClient(List<HostAddress> addresses) {
        this.connection = new GraphStorageConnection();
        this.addresses = addresses;
    }
 
    public StorageClient(List<HostAddress> addresses, int timeout) {
        this.connection = new GraphStorageConnection();
        this.addresses = addresses;
        this.timeout = timeout;
    }
 
    public StorageClient(List<HostAddress> addresses, int timeout, int connectionRetry,
                         int executionRetry, boolean enableSSL, SSLParam sslParam) {
        this(addresses, timeout);
        this.connectionRetry = connectionRetry;
        this.executionRetry = executionRetry;
        this.enableSSL = enableSSL;
        this.sslParam = sslParam;
        if (enableSSL && sslParam == null) {
            throw new IllegalArgumentException("SSL is enabled, but SSLParam is nul.");
        }
    }
 
    public boolean connect() throws Exception {
        connection.open(addresses.get(0), timeout, enableSSL, sslParam);
        StoragePoolConfig config = new StoragePoolConfig();
        config.setEnableSSL(enableSSL);
        config.setSslParam(sslParam);
        pool = new StorageConnPool(config);
        metaManager = new MetaManager(addresses, timeout, connectionRetry, executionRetry,
                enableSSL, sslParam);
        metaManager.addStorageAddrMapping(storageAddressMapping);
        return true;
    }
 
    public StorageClient setUser(String user) {
        this.user = user;
        return this;
    }
 
    public StorageClient setPassword(String password) {
        this.password = password;
        return this;
    }
 
 
    public void setStorageAddressMapping(Map<String, String> storageAddressMapping) {
        this.storageAddressMapping = storageAddressMapping;
        if (this.metaManager != null) {
            this.metaManager.addStorageAddrMapping(storageAddressMapping);
        }
    }
 
 
    public ScanVertexResultIterator scanVertex(String spaceName, String tagName,
                                               List<String> returnCols) {
        return scanVertex(spaceName, tagName, returnCols, DEFAULT_LIMIT);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName, int part, String tagName,
                                               List<String> returnCols) {
        return scanVertex(spaceName, part, tagName, returnCols, DEFAULT_LIMIT);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName, String tagName) {
        return scanVertex(spaceName, tagName, DEFAULT_LIMIT);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName, int part, String tagName) {
        return scanVertex(spaceName, part, tagName, DEFAULT_LIMIT);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit) {
        return scanVertex(spaceName, tagName, returnCols, limit, DEFAULT_START_TIME,
                DEFAULT_END_TIME);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit) {
        return scanVertex(spaceName, part, tagName, returnCols, limit, DEFAULT_START_TIME,
                DEFAULT_END_TIME);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               int limit) {
        return scanVertex(spaceName, tagName, limit, DEFAULT_START_TIME, DEFAULT_END_TIME);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               int limit) {
        return scanVertex(spaceName, part, tagName,
                limit, DEFAULT_START_TIME, DEFAULT_END_TIME);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit,
                                               long startTime,
                                               long endTime) {
        return scanVertex(spaceName, tagName, returnCols, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit,
                                               long startTime,
                                               long endTime) {
        return scanVertex(spaceName, part, tagName, returnCols, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               int limit,
                                               long startTime,
                                               long endTime) {
        return scanVertex(spaceName, tagName, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               int limit,
                                               long startTime,
                                               long endTime) {
        return scanVertex(spaceName,
                part,
                tagName,
                new ArrayList<>(),
                limit,
                startTime,
                endTime,
                DEFAULT_ALLOW_PART_SUCCESS,
                DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit,
                                               long startTime,
                                               long endTime,
                                               boolean allowPartSuccess,
                                               boolean allowReadFromFollower) {
        List<Integer> parts = metaManager.getSpaceParts(spaceName);
        if (parts.isEmpty()) {
            throw new IllegalArgumentException("No valid part in space " + spaceName);
        }
        return scanVertex(
                spaceName,
                parts,
                tagName,
                returnCols,
                false,
                limit,
                startTime,
                endTime,
                allowPartSuccess,
                allowReadFromFollower);
    }
 
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               List<String> returnCols,
                                               int limit,
                                               long startTime,
                                               long endTime,
                                               boolean allowPartSuccess,
                                               boolean allowReadFromFollower) {
        return scanVertex(
                spaceName,
                Arrays.asList(part),
                tagName,
                returnCols,
                false,
                limit,
                startTime,
                endTime,
                allowPartSuccess,
                allowReadFromFollower);
    }
 
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               String tagName,
                                               int limit,
                                               long startTime,
                                               long endTime,
                                               boolean allowPartSuccess,
                                               boolean allowReadFromFollower) {
        List<Integer> parts = metaManager.getSpaceParts(spaceName);
        if (parts.isEmpty()) {
            throw new IllegalArgumentException("No valid part in space " + spaceName);
        }
        return scanVertex(
                spaceName,
                parts,
                tagName,
                new ArrayList<>(),
                true,
                limit,
                startTime,
                endTime,
                allowPartSuccess,
                allowReadFromFollower);
    }
 
    public ScanVertexResultIterator scanVertex(String spaceName,
                                               int part,
                                               String tagName,
                                               int limit,
                                               long startTime,
                                               long endTime,
                                               boolean allowPartSuccess,
                                               boolean allowReadFromFollower) {
        return scanVertex(
                spaceName,
                Arrays.asList(part),
                tagName,
                new ArrayList<>(),
                true,
                limit,
                startTime,
                endTime,
                allowPartSuccess,
                allowReadFromFollower);
    }
 
    private ScanVertexResultIterator scanVertex(String spaceName,
                                                List<Integer> parts,
                                                String tagName,
                                                List<String> returnCols,
                                                boolean noColumns,
                                                int limit,
                                                long startTime,
                                                long endTime,
                                                boolean allowPartSuccess,
                                                boolean allowReadFromFollower) {
        if (spaceName == null || spaceName.trim().isEmpty()) {
            throw new IllegalArgumentException("space name is empty.");
        }
        if (tagName == null || tagName.trim().isEmpty()) {
            throw new IllegalArgumentException("tag name is empty");
        }
        if (noColumns && returnCols == null) {
            throw new IllegalArgumentException("returnCols is null");
        }
 
        Set<PartScanInfo> partScanInfoSet = new HashSet<>();
        for (int part : parts) {
            HostAddr leader = metaManager.getLeader(spaceName, part);
            partScanInfoSet.add(new PartScanInfo(part, new HostAddress(leader.getHost(),
                    leader.getPort())));
        }
        List<HostAddress> addrs = new ArrayList<>();
        for (HostAddr addr : metaManager.listHosts()) {
            addrs.add(new HostAddress(addr.getHost(), addr.getPort()));
        }
 
        long tag = metaManager.getTag(spaceName, tagName).getTag_id();
        List<byte[]> props = new ArrayList<>();
        props.add("_vid".getBytes());
        if (!noColumns) {
            if (returnCols.size() == 0) {
                Schema schema = metaManager.getTag(spaceName, tagName).getSchema();
                for (ColumnDef columnDef : schema.getColumns()) {
                    props.add(columnDef.getName());
                }
            } else {
                for (String prop : returnCols) {
                    props.add(prop.getBytes());
                }
            }
        }
        VertexProp vertexCols = new VertexProp((int) tag, props);
        List<VertexProp> vertexProps = Arrays.asList(vertexCols);
        ScanVertexRequest request = new ScanVertexRequest();
        request
                .setSpace_id(getSpaceId(spaceName))
                .setReturn_columns(vertexProps)
                .setLimit(limit)
                .setStart_time(startTime)
                .setEnd_time(endTime)
                .setEnable_read_from_follower(allowReadFromFollower);
 
        return doScanVertex(spaceName, tagName, partScanInfoSet, request, addrs, allowPartSuccess);
    }
 
 
    private ScanVertexResultIterator doScanVertex(String spaceName,
                                                  String tagName,
                                                  Set<PartScanInfo> partScanInfoSet,
                                                  ScanVertexRequest request,
                                                  List<HostAddress> addrs,
                                                  boolean allowPartSuccess) {
        if (addrs == null || addrs.isEmpty()) {
            throw new IllegalArgumentException("storage hosts is empty.");
        }
 
        return new ScanVertexResultIterator.ScanVertexResultBuilder()
                .withMetaClient(metaManager)
                .withPool(pool)
                .withPartScanInfo(partScanInfoSet)
                .withRequest(request)
                .withAddresses(addrs)
                .withSpaceName(spaceName)
                .withTagName(tagName)
                .withPartSuccess(allowPartSuccess)
                .withUser(user)
                .withPassword(password)
                .withStorageAddressMapping(storageAddressMapping)
                .build();
    }
 
 
    public ScanEdgeResultIterator scanEdge(String spaceName, String edgeName,
                                           List<String> returnCols) {
 
        return scanEdge(spaceName, edgeName, returnCols, DEFAULT_LIMIT);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, int part, String edgeName,
                                           List<String> returnCols) {
 
        return scanEdge(spaceName, part, edgeName, returnCols, DEFAULT_LIMIT);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, String edgeName) {
        return scanEdge(spaceName, edgeName, DEFAULT_LIMIT);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, int part, String edgeName) {
        return scanEdge(spaceName, part, edgeName, DEFAULT_LIMIT);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, String edgeName,
                                           List<String> returnCols, int limit) {
        return scanEdge(spaceName, edgeName, returnCols, limit, DEFAULT_START_TIME,
                DEFAULT_END_TIME);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, int part, String edgeName,
                                           List<String> returnCols, int limit) {
        return scanEdge(spaceName, part, edgeName, returnCols, limit, DEFAULT_START_TIME,
                DEFAULT_END_TIME);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, String edgeName, int limit) {
        return scanEdge(spaceName, edgeName, limit, DEFAULT_START_TIME, DEFAULT_END_TIME);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName, int part, String edgeName, int limit) {
        return scanEdge(spaceName, part, edgeName, limit, DEFAULT_START_TIME, DEFAULT_END_TIME);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           String edgeName,
                                           List<String> returnCols,
                                           int limit,
                                           long startTime,
                                           long endTime) {
        return scanEdge(spaceName, edgeName, returnCols, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           int part,
                                           String edgeName,
                                           List<String> returnCols,
                                           int limit,
                                           long startTime,
                                           long endTime) {
        return scanEdge(spaceName, part, edgeName, returnCols, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           String edgeName,
                                           int limit,
                                           long startTime,
                                           long endTime) {
        return scanEdge(spaceName, edgeName, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           int part,
                                           String edgeName,
                                           int limit,
                                           long startTime,
                                           long endTime) {
        return scanEdge(spaceName, part, edgeName, limit, startTime, endTime,
                DEFAULT_ALLOW_PART_SUCCESS, DEFAULT_ALLOW_READ_FOLLOWER);
    }
 
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           String edgeName,
                                           List<String> returnCols,
                                           int limit,
                                           long startTime,
                                           long endTime,
                                           boolean allowPartSuccess,
                                           boolean allowReadFromFollower) {
 
        List<Integer> parts = metaManager.getSpaceParts(spaceName);
        if (parts.isEmpty()) {
            throw new IllegalArgumentException("No valid part in space " + spaceName);
        }
 
        return scanEdge(spaceName, parts, edgeName, returnCols, false,
                limit, startTime, endTime, allowPartSuccess, allowReadFromFollower);
    }
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           int part,
                                           String edgeName,
                                           List<String> returnCols,
                                           int limit,
                                           long startTime,
                                           long endTime,
                                           boolean allowPartSuccess,
                                           boolean allowReadFromFollower) {
        return scanEdge(spaceName, Arrays.asList(part), edgeName, returnCols, false,
                limit, startTime, endTime, allowPartSuccess, allowReadFromFollower);
    }
 
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           String edgeName,
                                           int limit,
                                           long startTime,
                                           long endTime,
                                           boolean allowPartSuccess,
                                           boolean allowReadFromFollower) {
 
        List<Integer> parts = metaManager.getSpaceParts(spaceName);
        if (parts.isEmpty()) {
            throw new IllegalArgumentException("No valid part in space " + spaceName);
        }
        return scanEdge(spaceName, parts, edgeName, new ArrayList<>(), true,
                limit, startTime, endTime, allowPartSuccess, allowReadFromFollower);
    }
 
 
    public ScanEdgeResultIterator scanEdge(String spaceName,
                                           int part,
                                           String edgeName,
                                           int limit,
                                           long startTime,
                                           long endTime,
                                           boolean allowPartSuccess,
                                           boolean allowReadFromFollower) {
        return scanEdge(spaceName, Arrays.asList(part), edgeName, new ArrayList<>(), true,
                limit, startTime, endTime, allowPartSuccess, allowReadFromFollower);
    }
 
 
    private ScanEdgeResultIterator scanEdge(String spaceName,
                                            List<Integer> parts,
                                            String edgeName,
                                            List<String> returnCols,
                                            boolean noColumns,
                                            int limit,
                                            long startTime,
                                            long endTime,
                                            boolean allowPartSuccess,
                                            boolean allowReadFromFollower) {
        if (spaceName == null || spaceName.trim().isEmpty()) {
            throw new IllegalArgumentException("space name is empty.");
        }
        if (edgeName == null || edgeName.trim().isEmpty()) {
            throw new IllegalArgumentException("edge name is empty");
        }
        if (noColumns && returnCols == null) {
            throw new IllegalArgumentException("returnCols is null");
        }
 
        Set<PartScanInfo> partScanInfoSet = new HashSet<>();
        for (int part : parts) {
            HostAddr leader = metaManager.getLeader(spaceName, part);
            partScanInfoSet.add(new PartScanInfo(part, new HostAddress(leader.getHost(),
                    leader.getPort())));
        }
        List<HostAddress> addrs = new ArrayList<>();
        for (HostAddr addr : metaManager.listHosts()) {
            addrs.add(new HostAddress(addr.getHost(), addr.getPort()));
        }
        List<byte[]> props = new ArrayList<>();
        props.add("_src".getBytes());
        props.add("_dst".getBytes());
        props.add("_rank".getBytes());
        if (!noColumns) {
            if (returnCols.size() == 0) {
                Schema schema = metaManager.getEdge(spaceName, edgeName).getSchema();
                for (ColumnDef columnDef : schema.getColumns()) {
                    props.add(columnDef.name);
                }
            } else {
                for (String prop : returnCols) {
                    props.add(prop.getBytes());
                }
            }
        }
 
        long edgeId = getEdgeId(spaceName, edgeName);
        EdgeProp edgeCols = new EdgeProp((int) edgeId, props);
        List<EdgeProp> edgeProps = Arrays.asList(edgeCols);
 
        ScanEdgeRequest request = new ScanEdgeRequest();
        request
                .setSpace_id(getSpaceId(spaceName))
                .setReturn_columns(edgeProps)
                .setLimit(limit)
                .setStart_time(startTime)
                .setEnd_time(endTime)
                .setEnable_read_from_follower(allowReadFromFollower);
 
        return doScanEdge(spaceName, edgeName, partScanInfoSet, request, addrs, allowPartSuccess);
    }
 
 
    private ScanEdgeResultIterator doScanEdge(String spaceName,
                                              String edgeName,
                                              Set<PartScanInfo> partScanInfoSet,
                                              ScanEdgeRequest request,
                                              List<HostAddress> addrs,
                                              boolean allowPartSuccess) {
        if (addrs == null || addrs.isEmpty()) {
            throw new IllegalArgumentException("storage hosts is empty.");
        }
 
        return new ScanEdgeResultIterator.ScanEdgeResultBuilder()
                .withMetaClient(metaManager)
                .withPool(pool)
                .withPartScanInfo(partScanInfoSet)
                .withRequest(request)
                .withAddresses(addrs)
                .withSpaceName(spaceName)
                .withEdgeName(edgeName)
                .withPartSuccess(allowPartSuccess)
                .withUser(user)
                .withPassword(password)
                .build();
    }
 
 
    public void close() {
        if (pool != null) {
            pool.close();
        }
        if (connection != null) {
            connection.close();
        }
        if (metaManager != null) {
            metaManager.close();
        }
    }
 
 
    protected GraphStorageConnection getConnection() {
        return this.connection;
    }
 
 
    private int getSpaceId(String spaceName) {
        return metaManager.getSpaceId(spaceName);
    }
 
    private long getEdgeId(String spaceName, String edgeName) {
        return metaManager.getEdge(spaceName, edgeName).getEdge_type();
    }
 
    private static final int DEFAULT_LIMIT = 1000;
    private static final long DEFAULT_START_TIME = 0;
    private static final long DEFAULT_END_TIME = Long.MAX_VALUE;
    private static final boolean DEFAULT_ALLOW_PART_SUCCESS = false;
    private static final boolean DEFAULT_ALLOW_READ_FOLLOWER = false;
}