src/main/java/com/baidu/hugegraph/loader/parser/ElementParser.java

/*
 * Copyright 2017 HugeGraph Authors
 *
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with this
 * work for additional information regarding copyright ownership. The ASF
 * licenses this file to You under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * License for the specific language governing permissions and limitations
 * under the License.
 */

package com.baidu.hugegraph.loader.parser;

import java.io.UnsupportedEncodingException;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.apache.commons.lang3.StringUtils;

import com.baidu.hugegraph.driver.HugeClient;
import com.baidu.hugegraph.loader.exception.ParseException;
import com.baidu.hugegraph.loader.executor.LoadOptions;
import com.baidu.hugegraph.loader.reader.InputReader;
import com.baidu.hugegraph.loader.source.ElementSource;
import com.baidu.hugegraph.loader.util.DataTypeUtil;
import com.baidu.hugegraph.loader.util.HugeClientWrapper;
import com.baidu.hugegraph.structure.GraphElement;
import com.baidu.hugegraph.structure.SchemaElement;
import com.baidu.hugegraph.structure.constant.HugeType;
import com.baidu.hugegraph.structure.constant.IdStrategy;
import com.baidu.hugegraph.structure.schema.EdgeLabel;
import com.baidu.hugegraph.structure.schema.PropertyKey;
import com.baidu.hugegraph.structure.schema.VertexLabel;
import com.baidu.hugegraph.util.E;
import com.google.common.collect.HashBasedTable;
import com.google.common.collect.Table;

public abstract class ElementParser<GE extends GraphElement>
       implements Iterator<GE> {

    private static final int VERTEX_ID_LIMIT = 128;
    private static final String ID_CHARSET = "UTF-8";

    private final InputReader reader;

    private final HugeClient client;
    private final Table<HugeType, String, SchemaElement> schemas;

    ElementParser(InputReader reader, LoadOptions options) {
        this.reader = reader;
        this.client = HugeClientWrapper.get(options);
        this.schemas = HashBasedTable.create();
        this.reader.init();
    }

    public abstract ElementSource source();

    public InputReader reader() {
        return this.reader;
    }

    @Override
    public boolean hasNext() {
        return this.reader.hasNext();
    }

    @Override
    public GE next() {
        String line = this.reader().line();
        try {
            return this.parse(this.filterFields(this.reader().next()));
        } catch (IllegalArgumentException e) {
            throw new ParseException(line, e.getMessage());
        }
    }

    protected abstract GE parse(Map<String, Object> keyValues);

    protected abstract boolean isIdField(String fieldName);

    protected Map<String, Object> filterFields(Map<String, Object> keyValues) {
        for (String field : this.source().ignoredFields()) {
            keyValues.remove(field);
        }
        return keyValues;
    }

    protected void addProperties(GE element, Map<String, Object> keyValues) {
        for (Map.Entry<String, Object> entry : keyValues.entrySet()) {
            String fieldName = entry.getKey();
            Object fieldValue = entry.getValue();
            if (this.isIdField(fieldName)) {
                continue;
            }
            String key = this.source().mappingField(fieldName);
            Object value = this.validatePropertyValue(key, fieldValue);

            element.property(key, value);
        }
    }

    protected PropertyKey getPropertyKey(String name) {
        SchemaElement schema = this.schemas.get(HugeType.PROPERTY_KEY, name);
        if (schema == null) {
            schema = this.client.schema().getPropertyKey(name);
        }
        if (schema == null) {
            throw new IllegalStateException(
                      String.format("The property key %s doesn't exist", name));
        } else {
            this.schemas.put(HugeType.PROPERTY_KEY, name, schema);
        }
        return (PropertyKey) schema;
    }

    protected VertexLabel getVertexLabel(String name) {
        SchemaElement schema = this.schemas.get(HugeType.VERTEX_LABEL, name);
        if (schema == null) {
            schema = this.client.schema().getVertexLabel(name);
        }
        if (schema == null) {
            throw new IllegalStateException(
                      String.format("The vertex label %s doesn't exist", name));
        } else {
            this.schemas.put(HugeType.VERTEX_LABEL, name, schema);
        }
        return (VertexLabel) schema;
    }

    protected EdgeLabel getEdgeLabel(String name) {
        SchemaElement schema = this.schemas.get(HugeType.EDGE_LABEL, name);
        if (schema == null) {
            schema = this.client.schema().getEdgeLabel(name);
        }
        if (schema == null) {
            throw new IllegalStateException(
                      String.format("The edge label %s doesn't exist", name));
        } else {
            this.schemas.put(HugeType.EDGE_LABEL, name, schema);
        }
        return (EdgeLabel) schema;
    }

    protected String spliceVertexId(VertexLabel vertexLabel,
                                    Object[] primaryValues) {
        E.checkArgument(vertexLabel.primaryKeys().size() == primaryValues.length,
                        "Missing some primary key columns, expect %s, " +
                        "but only got %s for vertex label '%s'",
                        vertexLabel.primaryKeys(), primaryValues, vertexLabel);

        StringBuilder vertexId = new StringBuilder();
        StringBuilder vertexKeysId = new StringBuilder();
        String[] searchList = new String[]{":", "!"};
        String[] replaceList = new String[]{"`:", "`!"};
        for (Object value : primaryValues) {
            String pkValue = String.valueOf(value);
            pkValue = StringUtils.replaceEach(pkValue, searchList, replaceList);
            vertexKeysId.append(pkValue);
            vertexKeysId.append("!");
        }

        vertexId.append(vertexLabel.id()).append(":").append(vertexKeysId);
        vertexId.deleteCharAt(vertexId.length() - 1);
        return vertexId.toString();
    }

    protected void checkVertexIdLength(String id) {
        try {
            E.checkArgument(id.getBytes(ID_CHARSET).length <= VERTEX_ID_LIMIT,
                            "Vertex id length limit is '%s', '%s' exceeds it",
                            VERTEX_ID_LIMIT, id);
        } catch (UnsupportedEncodingException e) {
            throw new IllegalStateException(e);
        }
    }

    protected Object validatePropertyValue(String key, Object rawValue) {
        PropertyKey pKey = this.getPropertyKey(key);
        Object value = DataTypeUtil.convert(rawValue, pKey);
        E.checkArgument(value != null,
                        "The value '%s' can't convert to class %s " +
                        "with cardinality %s",
                        rawValue, pKey.dataType().clazz(), pKey.cardinality());
        return value;
    }

    public static boolean isAutomatic(IdStrategy idStrategy) {
        return idStrategy == IdStrategy.AUTOMATIC;
    }

    public static boolean isCustomize(IdStrategy idStrategy) {
        return idStrategy == IdStrategy.CUSTOMIZE_STRING ||
               idStrategy == IdStrategy.CUSTOMIZE_NUMBER;
    }

    public static boolean isPrimaryKey(IdStrategy idStrategy) {
        return idStrategy == IdStrategy.PRIMARY_KEY;
    }
}