ATLAS-1510: Consolidate/batch calls to GraphBackedTypeStore.findVertex()

release-log.txt

@@ -9,6 +9,7 @@ ATLAS-1060 Add composite indexes for exact match performance improvements for al
 ATLAS-1127 Modify creation and modification timestamps to Date instead of Long(sumasai)
 
 ALL CHANGES:
+ATLAS-1510 Consolidate/batch calls to GraphBackedTypeStore.findVertex() (jnhagelb)
 ATLAS-1388 Cache entities that are created/updated (jnhagelb)
 ATLAS-1369 Optimize Gremlin queries generated by DSL translator (jnhagelb)
 ATLAS-1517: updated hive_model to include schema related attributes (sarath.kum4r@gmail.com via mneethiraj)

repository/src/main/java/org/apache/atlas/repository/graph/GraphHelper.java

@@ -538,7 +538,7 @@ public final class GraphHelper {
      *
      *  @return propertyValue to AtlasVertex map with the result.
      */
-    private Map<String, AtlasVertex> getVerticesForPropertyValues(String property, List<String> values)
+    public Map<String, AtlasVertex> getVerticesForPropertyValues(String property, List<String> values)
             throws RepositoryException {
 
         if(values.isEmpty()) {

repository/src/main/java/org/apache/atlas/repository/typestore/TypePersistenceVisitor.java

+/**
+ * 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 org.apache.atlas.repository.typestore;
+
+import static org.apache.atlas.repository.graph.GraphHelper.setProperty;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.atlas.AtlasException;
+import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.RepositoryException;
+import org.apache.atlas.repository.graph.GraphHelper;
+import org.apache.atlas.repository.graphdb.AtlasVertex;
+import org.apache.atlas.typesystem.types.AttributeInfo;
+import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
+import org.apache.atlas.typesystem.types.EnumType;
+import org.apache.atlas.typesystem.types.EnumValue;
+import org.apache.atlas.typesystem.types.HierarchicalType;
+import org.apache.atlas.typesystem.types.IDataType;
+import org.apache.atlas.typesystem.types.TypeSystem;
+import org.codehaus.jettison.json.JSONException;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * TypeVisitor implementation that completes the type storage process by
+ * adding the required properties and edges to the type vertices
+ * that were created.
+ */
+public class TypePersistenceVisitor implements TypeVisitor {
+
+    private static final Logger LOG = LoggerFactory.getLogger(TypePersistenceVisitor.class);
+    private static final GraphHelper graphHelper = GraphHelper.getInstance();
+
+    private final GraphBackedTypeStore typeStore_;
+    private final Map<String,AtlasVertex> typeVertices;
+    private final TypeSystem typeSystem;
+
+    /**
+     * @param graphBackedTypeStore
+     */
+    public TypePersistenceVisitor(GraphBackedTypeStore graphBackedTypeStore, Map<String,AtlasVertex> typeVertices, TypeSystem typeSystem) {
+        typeStore_ = graphBackedTypeStore;
+        this.typeVertices = typeVertices;
+        this.typeSystem = typeSystem;
+    }
+
+    @Override
+    public void visitEnumeration(EnumType dataType) throws AtlasException {
+        AtlasVertex vertex = typeVertices.get(dataType.getName());
+        List<String> values = new ArrayList<>(dataType.values().size());
+        for (EnumValue enumValue : dataType.values()) {
+            String key = GraphBackedTypeStore.getPropertyKey(dataType.getName(), enumValue.value);
+            setProperty(vertex, key, enumValue.ordinal);
+            values.add(enumValue.value);
+        }
+        setProperty(vertex, GraphBackedTypeStore.getPropertyKey(dataType.getName()), values);
+
+    }
+    @Override
+    public void visitAttributeDataType(String typeName, AttributeInfo attribute, IDataType attrType) throws AtlasException {
+        AtlasVertex vertex = typeVertices.get(typeName);
+        String vertexTypeName = GraphHelper.getSingleValuedProperty(vertex, Constants.TYPENAME_PROPERTY_KEY, String.class);
+        AtlasVertex attrVertex = typeVertices.get(attrType.getName());
+        String label = GraphBackedTypeStore.getEdgeLabel(vertexTypeName, attribute.name);
+        graphHelper.getOrCreateEdge(vertex, attrVertex, label);
+    }
+    @Override
+    public void visitSuperType(String typeName, String superTypeName) throws AtlasException {
+        AtlasVertex vertex = typeVertices.get(typeName);
+        HierarchicalType superType = typeSystem.getDataType(HierarchicalType.class, superTypeName);
+        AtlasVertex superVertex = typeVertices.get(superTypeName);
+        graphHelper.getOrCreateEdge(vertex, superVertex, GraphBackedTypeStore.SUPERTYPE_EDGE_LABEL);
+    }
+
+    @Override
+    public void visitAttributeNames(String typeName, List<String> attrNames) throws AtlasException {
+        AtlasVertex vertex = typeVertices.get(typeName);
+        setProperty(vertex, GraphBackedTypeStore.getPropertyKey(typeName), attrNames);
+
+    }
+
+    @Override
+    public void visitAttribute(String typeName, AttributeInfo attribute) throws AtlasException {
+        AtlasVertex vertex = typeVertices.get(typeName);
+        String propertyKey = GraphBackedTypeStore.getPropertyKey(typeName, attribute.name);
+        try {
+            setProperty(vertex, propertyKey, attribute.toJson());
+        } catch (JSONException e) {
+            throw new StorageException(typeName, e);
+        }
+    }
+
+    @Override
+    public void visitDataType(TypeCategory category, String typeName, String typeDescription) {
+        //nothing to do
+
+    }
+}
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/typestore/TypeVertexFinder.java

+/**
+ * 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 org.apache.atlas.repository.typestore;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import org.apache.atlas.AtlasException;
+import org.apache.atlas.typesystem.types.AttributeInfo;
+import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
+import org.apache.atlas.typesystem.types.EnumType;
+import org.apache.atlas.typesystem.types.HierarchicalType;
+import org.apache.atlas.typesystem.types.IDataType;
+import org.apache.atlas.typesystem.types.TypeSystem;
+
+/**
+ * TypeVisitor implementation that builds up a list of type vertices
+ * that need to be created for the types that are being stored.
+ *
+ */
+public class TypeVertexFinder implements TypeVisitor {
+
+    private final List<TypeVertexInfo> toCreate = new ArrayList<TypeVertexInfo>();
+    private final Set<String> typesIncluded = new HashSet<String>();
+    private final TypeSystem typeSystem;
+
+    public TypeVertexFinder(TypeSystem ts) {
+        typeSystem = ts;
+    }
+
+
+    @Override
+    public void visitEnumeration(EnumType dataType) {
+        visitDataType(dataType);
+    }
+
+    private void addTypeIfNeeded(TypeVertexInfo info) {
+        if(! typesIncluded.contains(info.getTypeName())) {
+            toCreate.add(info);
+            typesIncluded.add(info.getTypeName());
+        }
+    }
+
+    @Override
+    public void visitAttributeDataType(String typeName, AttributeInfo sourceAttr, IDataType attrType) throws AtlasException {
+        visitDataType(attrType);
+    }
+
+    @Override
+    public void visitSuperType(String typeName, String superTypeName) throws AtlasException {
+        HierarchicalType superType = typeSystem.getDataType(HierarchicalType.class, superTypeName);
+        visitDataType(superType);
+    }
+
+    @Override
+    public void visitAttributeNames(String typeName, List<String> attrNames) throws AtlasException {
+       //nothing to do
+
+    }
+
+    @Override
+    public void visitAttribute(String typeName, AttributeInfo attribute) throws StorageException, AtlasException {
+        //nothing to do
+    }
+
+
+    private void visitDataType(IDataType dataType) {
+        TypeVertexInfo info = null;
+        info = new TypeVertexInfo(dataType.getTypeCategory(), dataType.getName(), dataType.getDescription());
+        addTypeIfNeeded(info);
+
+    }
+
+
+    public List<TypeVertexInfo> getVerticesToCreate() {
+        return toCreate;
+    }
+
+    @Override
+    public void visitDataType(TypeCategory category, String typeName, String typeDescription) {
+        TypeVertexInfo info = new TypeVertexInfo(category, typeName, typeDescription);
+        addTypeIfNeeded(info);
+    }
+
+}
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/typestore/TypeVertexInfo.java

+/**
+ * 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 org.apache.atlas.repository.typestore;
+
+import java.util.Objects;
+
+import org.apache.atlas.typesystem.types.DataTypes;
+import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
+
+/**
+ * Records the information needed to create a particular type vertex.
+ */
+public class TypeVertexInfo {
+
+    private DataTypes.TypeCategory category;
+    private String typeName;
+    private String typeDescription;
+
+    public TypeVertexInfo(TypeCategory category, String typeName, String typeDescription) {
+        super();
+        this.category = category;
+        this.typeName = typeName;
+        this.typeDescription = typeDescription;
+    }
+
+    public DataTypes.TypeCategory getCategory() {
+        return category;
+    }
+
+    public void setCategory(DataTypes.TypeCategory category) {
+        this.category = category;
+    }
+
+    public String getTypeName() {
+        return typeName;
+    }
+
+    public void setTypeName(String typeName) {
+        this.typeName = typeName;
+    }
+
+    public String getTypeDescription() {
+        return typeDescription;
+    }
+
+    public void setTypeDescription(String typeDescription) {
+        this.typeDescription = typeDescription;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(category, typeName);
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+
+        if (this == obj) {
+            return true;
+        }
+
+        if (getClass() != obj.getClass()) {
+            return false;
+        }
+
+        TypeVertexInfo other = (TypeVertexInfo)obj;
+        if(! Objects.equals(category, other.category)) {
+            return false;
+        }
+
+        if(! Objects.equals(typeName, other.typeName)) {
+            return false;
+        }
+
+        return true;
+    }
+
+}
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/typestore/TypeVisitor.java

+/**
+ * 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 org.apache.atlas.repository.typestore;
+
+import java.util.List;
+
+import org.apache.atlas.AtlasException;
+import org.apache.atlas.repository.RepositoryException;
+import org.apache.atlas.typesystem.types.AttributeInfo;
+import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
+import org.apache.atlas.typesystem.types.EnumType;
+import org.apache.atlas.typesystem.types.IDataType;
+
+/**
+ * Callback mechanism used when storing types.  As {@link GraphBackedTypeStore} traverses
+ * through the types being persisted, these methods are called with the information that
+ * it finds.
+ */
+public interface TypeVisitor {
+
+    /**
+     * Called when an enumeration type is found
+     * @param type
+     * @throws AtlasException
+     */
+    void visitEnumeration(EnumType type) throws AtlasException;
+
+    /**
+     * Called with a data type that is associated with a given attribute.  There can
+     * be more than one.  For example, map types have both a key and a value type.
+     * This is called once for each type.  This is called once for each datatype
+     * associated with the given attribute.
+     *
+     * @param typeName The name of the type being processed.
+     * @param sourceAttr The attribute in that type that we are processing.
+     * @param attrType A dataType associated with that attribute.
+     * @throws AtlasException
+     */
+    void visitAttributeDataType(String typeName, AttributeInfo sourceAttr, IDataType attrType) throws AtlasException;
+
+    /**
+     * Called when a super type is found.  It is called once for each superType.
+     *
+     * @param typeName The type being processed.
+     * @param superType The name of the super type that was found.
+     * @throws RepositoryException
+     * @throws AtlasException
+     */
+    void visitSuperType(String typeName, String superType) throws RepositoryException, AtlasException;
+
+    /**
+     * Called with the list of immediate attribute names that were found for the given type.  It
+     * is called once per type.
+     *
+     * @param typeName The name of the type that is being processed.
+     * @param attrNames The names of all of the immediate attributes in the type.
+     * @throws AtlasException
+     */
+    void visitAttributeNames(String typeName, List<String> attrNames) throws AtlasException;
+
+    /**
+     * Called once for each immediate attribute in a type.
+     * @param typeName The name of the type that is being procesed
+     * @param attribute The immediate attribute that was found
+     *
+     * @throws StorageException
+     * @throws AtlasException
+     */
+    void visitAttribute(String typeName, AttributeInfo attribute) throws StorageException, AtlasException;
+
+    /**
+     * Called once for each struct, class, and trait type that was found.  It is
+     * called when we start processing that type.
+     *
+     * @param category The category of the type
+     * @param typeName The name of the type
+     * @param typeDescription The description of the type.
+     */
+    void visitDataType(TypeCategory category, String typeName, String typeDescription);
+}
\ No newline at end of file

repository/src/test/java/org/apache/atlas/repository/typestore/GraphBackedTypeStoreTest.java

@@ -23,6 +23,7 @@ import static org.apache.atlas.typesystem.types.utils.TypesUtil.createOptionalAt
 import static org.apache.atlas.typesystem.types.utils.TypesUtil.createRequiredAttrDef;
 import static org.apache.atlas.typesystem.types.utils.TypesUtil.createStructTypeDef;
 
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -32,6 +33,7 @@ import javax.inject.Inject;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.RepositoryMetadataModule;
 import org.apache.atlas.TestUtils;
+import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
@@ -200,11 +202,11 @@ public class GraphBackedTypeStoreTest {
         verifyEdges();
     }
 
-    private void verifyEdges() {
+    private void verifyEdges() throws RepositoryException {
         // ATLAS-474: verify that type update did not write duplicate edges to the type store.
         if (typeStore instanceof GraphBackedTypeStore) {
             GraphBackedTypeStore gbTypeStore = (GraphBackedTypeStore) typeStore;
-            AtlasVertex typeVertex = gbTypeStore.findVertex(TypeCategory.CLASS, "Department");
+            AtlasVertex typeVertex = gbTypeStore.findVertices(Collections.singletonList("Department")).get("Department");
             int edgeCount = countOutgoingEdges(typeVertex, gbTypeStore.getEdgeLabel("Department", "employees"));
             Assert.assertEquals(edgeCount, 1, "Should only be 1 edge for employees attribute on Department type AtlasVertex");
         }