ATLAS-3737: Refactoring: Support multiple instances of AtlasGraph

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraph.java

@@ -319,8 +319,8 @@ public class AtlasJanusGraph implements AtlasGraph<AtlasJanusVertex, AtlasJanusE
         }
     }
 
-    private JanusGraph getGraph() {
-        return getGraphInstance();
+    public JanusGraph getGraph() {
+        return this.janusGraph;
     }
 
     @Override

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/query/NativeJanusGraphQuery.java

@@ -56,7 +56,7 @@ public class NativeJanusGraphQuery implements NativeTinkerpopGraphQuery<AtlasJan
     private JanusGraphQuery<?> query;
 
     public NativeJanusGraphQuery(AtlasJanusGraph graph) {
-        this.query = AtlasJanusGraphDatabase.getGraphInstance().query();
+        this.query = graph.getGraph().query();
         this.graph = graph;
     }
 

repository/src/main/java/org/apache/atlas/discovery/EntityDiscoveryService.java

@@ -99,7 +99,7 @@ public class EntityDiscoveryService implements AtlasDiscoveryService {
                            SearchTracker searchTracker,
                            UserProfileService userProfileService) throws AtlasException {
         this.graph                    = graph;
-        this.entityRetriever          = new EntityGraphRetriever(typeRegistry);
+        this.entityRetriever          = new EntityGraphRetriever(this.graph, typeRegistry);
         this.indexer                  = indexer;
         this.searchTracker            = searchTracker;
         this.gremlinQueryProvider     = AtlasGremlinQueryProvider.INSTANCE;

repository/src/main/java/org/apache/atlas/discovery/EntityLineageService.java

@@ -94,7 +94,7 @@ public class EntityLineageService implements AtlasLineageService {
     EntityLineageService(AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph) {
         this.graph = atlasGraph;
         this.gremlinQueryProvider = AtlasGremlinQueryProvider.INSTANCE;
-        this.entityRetriever = new EntityGraphRetriever(typeRegistry);
+        this.entityRetriever = new EntityGraphRetriever(atlasGraph, typeRegistry);
         this.atlasTypeRegistry = typeRegistry;
     }
 
@@ -259,7 +259,7 @@ public class EntityLineageService implements AtlasLineageService {
         }
 
         if (isDataSet) {
-            AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex datasetVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
             if (direction == INPUT || direction == BOTH) {
                 traverseEdges(datasetVertex, true, depth, ret);
@@ -269,7 +269,7 @@ public class EntityLineageService implements AtlasLineageService {
                 traverseEdges(datasetVertex, false, depth, ret);
             }
         } else  {
-            AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex processVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
             // make one hop to the next dataset vertices from process vertex and traverse with 'depth = depth - 1'
             if (direction == INPUT || direction == BOTH) {

repository/src/main/java/org/apache/atlas/repository/converters/AtlasInstanceConverter.java

@@ -34,6 +34,7 @@ import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.instance.EntityMutations.EntityOperation;
 import org.apache.atlas.model.instance.GuidMapping;
 import org.apache.atlas.model.legacy.EntityResult;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
 import org.apache.atlas.v1.model.instance.Referenceable;
 import org.apache.atlas.v1.model.instance.Struct;
@@ -68,11 +69,11 @@ public class AtlasInstanceConverter {
     private final EntityGraphRetriever  entityGraphRetrieverIgnoreRelationshipAttrs;
 
     @Inject
-    public AtlasInstanceConverter(AtlasTypeRegistry typeRegistry, AtlasFormatConverters instanceFormatters) {
+    public AtlasInstanceConverter(AtlasGraph graph, AtlasTypeRegistry typeRegistry, AtlasFormatConverters instanceFormatters) {
         this.typeRegistry                                = typeRegistry;
         this.instanceFormatters                          = instanceFormatters;
-        this.entityGraphRetriever                        = new EntityGraphRetriever(typeRegistry);
-        this.entityGraphRetrieverIgnoreRelationshipAttrs = new EntityGraphRetriever(typeRegistry, true);
+        this.entityGraphRetriever                        = new EntityGraphRetriever(graph, typeRegistry);
+        this.entityGraphRetrieverIgnoreRelationshipAttrs = new EntityGraphRetriever(graph, typeRegistry, true);
     }
 
     public Referenceable[] getReferenceables(Collection<AtlasEntity> entities) throws AtlasBaseException {

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

@@ -25,6 +25,7 @@ import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityExtInfo;
 import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
 import org.apache.atlas.type.AtlasArrayType;
 import org.apache.atlas.type.AtlasBuiltInTypes;
@@ -68,13 +69,13 @@ public class FullTextMapperV2 implements IFullTextMapper {
 
 
     @Inject
-    public FullTextMapperV2(AtlasTypeRegistry typeRegistry, Configuration configuration) {
+    public FullTextMapperV2(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, Configuration configuration) {
         this.typeRegistry  = typeRegistry;
         this.configuration = configuration;
 
         followReferences = this.configuration != null && this.configuration.getBoolean(FULL_TEXT_FOLLOW_REFERENCES, false);
         // If followReferences = false then ignore relationship attr loading
-        entityGraphRetriever = new EntityGraphRetriever(typeRegistry, !followReferences);
+        entityGraphRetriever = new EntityGraphRetriever(atlasGraph, typeRegistry, !followReferences);
     }
 
     /**

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

@@ -30,8 +30,11 @@ import org.apache.atlas.listener.ActiveStateChangeHandler;
 import org.apache.atlas.listener.ChangedTypeDefs;
 import org.apache.atlas.listener.TypeDefChangeListener;
 import org.apache.atlas.model.TypeCategory;
+import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
+import org.apache.atlas.model.typedef.AtlasEntityDef;
 import org.apache.atlas.model.typedef.AtlasEnumDef;
+import org.apache.atlas.model.typedef.AtlasRelationshipDef;
 import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.Constants;
@@ -41,8 +44,12 @@ import org.apache.atlas.repository.graphdb.*;
 import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
 import org.apache.atlas.type.*;
 import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
+import org.apache.atlas.type.AtlasType;
+import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.type.AtlasTypeUtil;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.configuration.Configuration;
+import org.apache.solr.common.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Component;
@@ -192,6 +199,10 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
 
             //resolve index fields names for the new entity attributes.
             resolveIndexFieldNames(management, changedTypeDefs);
+
+            createEdgeLabels(management, changedTypeDefs.getCreatedTypeDefs());
+            createEdgeLabels(management, changedTypeDefs.getUpdatedTypeDefs());
+
             //Commit indexes
             commit(management);
         } catch (RepositoryException | IndexException e) {
@@ -349,6 +360,13 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
             // create fulltext indexes
             createFullTextIndex(management, ENTITY_TEXT_PROPERTY_KEY, String.class, SINGLE);
 
+            createPropertyKey(management, IS_PROXY_KEY, Boolean.class, SINGLE);
+            createPropertyKey(management, PROVENANCE_TYPE_KEY, Integer.class, SINGLE);
+            createPropertyKey(management, HOME_ID_KEY, String.class, SINGLE);
+
+            createEdgeLabel(management, Constants.TERM_ASSIGNMENT_LABEL);
+            createEdgeLabel(management, Constants.CLASSIFICATION_LABEL);
+
             commit(management);
 
             LOG.info("Index creation for global keys complete.");
@@ -689,6 +707,14 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
 
         String label = Constants.INTERNAL_PROPERTY_KEY_PREFIX + propertyName;
 
+        createEdgeLabelUsingLabelName(management, label);
+    }
+
+    private void createEdgeLabelUsingLabelName(final AtlasGraphManagement management, final String label) {
+        if (StringUtils.isEmpty(label)) {
+            return;
+        }
+
         org.apache.atlas.repository.graphdb.AtlasEdgeLabel edgeLabel = management.getEdgeLabel(label);
 
         if (edgeLabel == null) {
@@ -999,4 +1025,47 @@ public class GraphBackedSearchIndexer implements SearchIndexer, ActiveStateChang
         }
     }
 
+
+    private void createEdgeLabels(AtlasGraphManagement management, List<? extends AtlasBaseTypeDef> typeDefs) {
+        if (CollectionUtils.isEmpty(typeDefs)) {
+            return;
+        }
+
+        for (AtlasBaseTypeDef typeDef : typeDefs) {
+            if (typeDef instanceof AtlasEntityDef) {
+                AtlasEntityDef entityDef = (AtlasEntityDef) typeDef;
+                createEdgeLabelsForStruct(management, entityDef);
+            } else if (typeDef instanceof AtlasRelationshipDef) {
+                createEdgeLabels(management, (AtlasRelationshipDef) typeDef);
+            }
+        }
+    }
+
+    private void createEdgeLabelsForStruct(AtlasGraphManagement management, AtlasEntityDef entityDef) {
+        try {
+            AtlasType type = typeRegistry.getType(entityDef.getName());
+            if (!(type instanceof AtlasEntityType)) {
+                return;
+            }
+
+            AtlasEntityType entityType = (AtlasEntityType) type;
+            for (AtlasAttributeDef attributeDef : entityDef.getAttributeDefs()) {
+                AtlasAttribute attribute = entityType.getAttribute(attributeDef.getName());
+                if (attribute.getAttributeType().getTypeCategory() == TypeCategory.STRUCT) {
+                    String relationshipLabel = attribute.getRelationshipEdgeLabel();
+                    createEdgeLabelUsingLabelName(management, relationshipLabel);
+                }
+            }
+        } catch (AtlasBaseException e) {
+            LOG.error("Error fetching type: {}", entityDef.getName(), e);
+        }
+    }
+
+    private void createEdgeLabels(AtlasGraphManagement management, AtlasRelationshipDef relationshipDef) {
+        String relationshipTypeName = relationshipDef.getName();
+        AtlasRelationshipType relationshipType = typeRegistry.getRelationshipTypeByName(relationshipTypeName);
+        String relationshipLabel = relationshipType.getRelationshipLabel();
+
+        createEdgeLabelUsingLabelName(management, relationshipLabel);
+    }
 }

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

@@ -95,18 +95,13 @@ public final class GraphHelper {
     public static final String RETRY_DELAY = "atlas.graph.storage.retry.sleeptime.ms";
     public static final String DEFAULT_REMOVE_PROPAGATIONS_ON_ENTITY_DELETE = "atlas.graph.remove.propagations.default";
 
-    private final AtlasGremlinQueryProvider queryProvider = AtlasGremlinQueryProvider.INSTANCE;
-
-    private static volatile GraphHelper INSTANCE;
-
     private AtlasGraph graph;
 
-    private static int     maxRetries;
-    private static long    retrySleepTimeMillis;
-    private static boolean removePropagations;
+    private int     maxRetries = 3;
+    private long    retrySleepTimeMillis = 1000;
+    private boolean removePropagations = false;
 
-    @VisibleForTesting
-    GraphHelper(AtlasGraph graph) {
+    public GraphHelper(AtlasGraph graph) {
         this.graph = graph;
         try {
             maxRetries           = ApplicationProperties.get().getInt(RETRY_COUNT, 3);
@@ -117,69 +112,10 @@ public final class GraphHelper {
         }
     }
 
-    public static GraphHelper getInstance() {
-        if ( INSTANCE == null) {
-            synchronized (GraphHelper.class) {
-                if (INSTANCE == null) {
-                    INSTANCE = new GraphHelper(AtlasGraphProvider.getGraphInstance());
-                }
-            }
-        }
-        return INSTANCE;
-    }
-
-    @VisibleForTesting
-    static GraphHelper getInstance(AtlasGraph graph) {
-        if ( INSTANCE == null) {
-            synchronized (GraphHelper.class) {
-                if (INSTANCE == null) {
-                    INSTANCE = new GraphHelper(graph);
-                }
-            }
-        }
-        return INSTANCE;
-    }
-
     public static boolean isTermEntityEdge(AtlasEdge edge) {
         return StringUtils.equals(edge.getLabel(), TERM_ASSIGNMENT_LABEL);
     }
 
-    public AtlasVertex createVertexWithIdentity(Referenceable typedInstance, Set<String> superTypeNames) {
-        final String guid = UUID.randomUUID().toString();
-
-        final AtlasVertex vertexWithIdentity = createVertexWithoutIdentity(typedInstance.getTypeName(),
-                new Id(guid, 0, typedInstance.getTypeName()), superTypeNames);
-
-        // add identity
-        AtlasGraphUtilsV2.setEncodedProperty(vertexWithIdentity, Constants.GUID_PROPERTY_KEY, guid);
-
-        // add version information
-        AtlasGraphUtilsV2.setEncodedProperty(vertexWithIdentity, Constants.VERSION_PROPERTY_KEY, Long.valueOf(typedInstance.getId().getVersion()));
-
-        return vertexWithIdentity;
-    }
-
-    public AtlasVertex createVertexWithoutIdentity(String typeName, Id typedInstanceId, Set<String> superTypeNames) {
-        if (LOG.isDebugEnabled()) {
-            LOG.debug("Creating AtlasVertex for type {} id {}", typeName, typedInstanceId != null ? typedInstanceId._getId() : null);
-        }
-
-        final AtlasVertex ret = graph.addVertex();
-
-        AtlasGraphUtilsV2.setEncodedProperty(ret, ENTITY_TYPE_PROPERTY_KEY, typeName);
-        AtlasGraphUtilsV2.setEncodedProperty(ret, STATE_PROPERTY_KEY, ACTIVE.name());
-        AtlasGraphUtilsV2.setEncodedProperty(ret, TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
-        AtlasGraphUtilsV2.setEncodedProperty(ret, MODIFICATION_TIMESTAMP_PROPERTY_KEY, RequestContext.get().getRequestTime());
-        AtlasGraphUtilsV2.setEncodedProperty(ret, CREATED_BY_KEY, RequestContext.get().getUser());
-        AtlasGraphUtilsV2.setEncodedProperty(ret, MODIFIED_BY_KEY, RequestContext.get().getUser());
-
-        for (String superTypeName : superTypeNames) {
-            AtlasGraphUtilsV2.addEncodedProperty(ret, SUPER_TYPES_PROPERTY_KEY, superTypeName);
-        }
-
-        return ret;
-    }
-
     public AtlasEdge addClassificationEdge(AtlasVertex entityVertex, AtlasVertex classificationVertex, boolean isPropagated) {
         AtlasEdge ret = addEdge(entityVertex, classificationVertex, CLASSIFICATION_LABEL);
 
@@ -1193,6 +1129,14 @@ public final class GraphHelper {
         return ret;
     }
 
+    public AtlasGraph getGraph() {
+        return this.graph;
+    }
+
+    public Boolean getDefaultRemovePropagations() {
+        return this.removePropagations;
+    }
+
     /**
      * Guid and AtlasVertex combo
      */
@@ -1663,10 +1607,6 @@ public final class GraphHelper {
         return StringUtils.equals(getGuid(vertexB), getGuid(vertexA));
     }
 
-    public static boolean getDefaultRemovePropagations() {
-        return removePropagations;
-    }
-
     public static String getDelimitedClassificationNames(Collection<String> classificationNames) {
         String ret = null;
 

repository/src/main/java/org/apache/atlas/repository/impexp/ExportService.java

@@ -66,14 +66,14 @@ public class ExportService {
     private final HdfsPathEntityCreator     hdfsPathEntityCreator;
 
     @Inject
-    public ExportService(final AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph,
+    public ExportService(final AtlasTypeRegistry typeRegistry, AtlasGraph graph,
                          AuditsWriter auditsWriter, HdfsPathEntityCreator hdfsPathEntityCreator) {
         this.typeRegistry         = typeRegistry;
-        this.entityGraphRetriever = new EntityGraphRetriever(this.typeRegistry);
+        this.entityGraphRetriever = new EntityGraphRetriever(graph, this.typeRegistry);
         this.auditsWriter         = auditsWriter;
         this.hdfsPathEntityCreator = hdfsPathEntityCreator;
-        this.startEntityFetchByExportRequest = new StartEntityFetchByExportRequest(atlasGraph, typeRegistry, AtlasGremlinQueryProvider.INSTANCE);
-        this.entitiesExtractor = new EntitiesExtractor(atlasGraph, typeRegistry);
+        this.startEntityFetchByExportRequest = new StartEntityFetchByExportRequest(graph, typeRegistry, AtlasGremlinQueryProvider.INSTANCE);
+        this.entitiesExtractor = new EntitiesExtractor(graph, typeRegistry);
     }
 
     public AtlasExportResult run(ZipSink exportSink, AtlasExportRequest request, String userName, String hostName,

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerDelegate.java

@@ -19,6 +19,7 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.RequestContext;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.store.DeleteType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.util.AtlasRepositoryConfiguration;
@@ -37,11 +38,13 @@ public class DeleteHandlerDelegate {
     private final SoftDeleteHandlerV1 softDeleteHandler;
     private final HardDeleteHandlerV1 hardDeleteHandler;
     private final DeleteHandlerV1     defaultHandler;
+    private final AtlasGraph atlasGraph;
 
     @Inject
-    public DeleteHandlerDelegate(AtlasTypeRegistry typeRegistry) {
-        this.softDeleteHandler = new SoftDeleteHandlerV1(typeRegistry);
-        this.hardDeleteHandler = new HardDeleteHandlerV1(typeRegistry);
+    public DeleteHandlerDelegate(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry) {
+        this.atlasGraph = atlasGraph;
+        this.softDeleteHandler = new SoftDeleteHandlerV1(atlasGraph, typeRegistry);
+        this.hardDeleteHandler = new HardDeleteHandlerV1(atlasGraph, typeRegistry);
         this.defaultHandler    = getDefaultConfiguredHandler(typeRegistry);
     }
 
@@ -74,7 +77,7 @@ public class DeleteHandlerDelegate {
 
             LOG.info("Default delete handler set to: {}", handlerFromProperties.getName());
 
-            ret = (DeleteHandlerV1) handlerFromProperties.getConstructor(AtlasTypeRegistry.class).newInstance(typeRegistry);
+            ret = (DeleteHandlerV1) handlerFromProperties.getConstructor(AtlasGraph.class, AtlasTypeRegistry.class).newInstance(this.atlasGraph, typeRegistry);
         } catch (Exception ex) {
             LOG.error("Error instantiating default delete handler. Defaulting to: {}", softDeleteHandler.getClass().getName(), ex);
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/DeleteHandlerV1.java

@@ -32,6 +32,7 @@ import org.apache.atlas.repository.graph.AtlasEdgeLabel;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
@@ -62,15 +63,17 @@ import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelation
 
 public abstract class DeleteHandlerV1 {
     public static final Logger LOG = LoggerFactory.getLogger(DeleteHandlerV1.class);
-    protected static final GraphHelper graphHelper = GraphHelper.getInstance();
+
+    protected final GraphHelper graphHelper;
     private final AtlasTypeRegistry    typeRegistry;
     private final EntityGraphRetriever entityRetriever;
     private final boolean              shouldUpdateInverseReferences;
     private final boolean              softDelete;
 
-    public DeleteHandlerV1(AtlasTypeRegistry typeRegistry, boolean shouldUpdateInverseReference, boolean softDelete) {
+    public DeleteHandlerV1(AtlasGraph graph, AtlasTypeRegistry typeRegistry, boolean shouldUpdateInverseReference, boolean softDelete) {
         this.typeRegistry                  = typeRegistry;
-        this.entityRetriever               = new EntityGraphRetriever(typeRegistry);
+        this.graphHelper                   = new GraphHelper(graph);
+        this.entityRetriever               = new EntityGraphRetriever(graph, typeRegistry);
         this.shouldUpdateInverseReferences = shouldUpdateInverseReference;
         this.softDelete                    = softDelete;
     }
@@ -210,7 +213,7 @@ public abstract class DeleteHandlerV1 {
                     if (attributeInfo.getAttributeDef().isSoftReferenced()) {
                         String        softRefVal = vertex.getProperty(attributeInfo.getVertexPropertyName(), String.class);
                         AtlasObjectId refObjId   = AtlasEntityUtil.parseSoftRefValue(softRefVal);
-                        AtlasVertex   refVertex  = refObjId != null ? AtlasGraphUtilsV2.findByGuid(refObjId.getGuid()) : null;
+                        AtlasVertex   refVertex  = refObjId != null ? AtlasGraphUtilsV2.findByGuid(this.graphHelper.getGraph(), refObjId.getGuid()) : null;
 
                         if (refVertex != null) {
                             vertices.push(refVertex);
@@ -244,7 +247,7 @@ public abstract class DeleteHandlerV1 {
 
                             if (CollectionUtils.isNotEmpty(refObjIds)) {
                                 for (AtlasObjectId refObjId : refObjIds) {
-                                    AtlasVertex refVertex = AtlasGraphUtilsV2.findByGuid(refObjId.getGuid());
+                                    AtlasVertex refVertex = AtlasGraphUtilsV2.findByGuid(this.graphHelper.getGraph(), refObjId.getGuid());
 
                                     if (refVertex != null) {
                                         vertices.push(refVertex);
@@ -257,7 +260,7 @@ public abstract class DeleteHandlerV1 {
 
                             if (MapUtils.isNotEmpty(refObjIds)) {
                                 for (AtlasObjectId refObjId : refObjIds.values()) {
-                                    AtlasVertex refVertex = AtlasGraphUtilsV2.findByGuid(refObjId.getGuid());
+                                    AtlasVertex refVertex = AtlasGraphUtilsV2.findByGuid(this.graphHelper.getGraph(), refObjId.getGuid());
 
                                     if (refVertex != null) {
                                         vertices.push(refVertex);

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/HardDeleteHandlerV1.java

@@ -22,6 +22,7 @@ import org.apache.atlas.annotation.ConditionalOnAtlasProperty;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.springframework.stereotype.Component;
@@ -33,8 +34,8 @@ import javax.inject.Inject;
 public class HardDeleteHandlerV1 extends DeleteHandlerV1 {
 
     @Inject
-    public HardDeleteHandlerV1(AtlasTypeRegistry typeRegistry) {
-        super(typeRegistry, true, false);
+    public HardDeleteHandlerV1(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        super(graph, typeRegistry, true, false);
     }
 
     @Override

repository/src/main/java/org/apache/atlas/repository/store/graph/v1/SoftDeleteHandlerV1.java

@@ -22,6 +22,7 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
 import org.apache.atlas.type.AtlasTypeRegistry;
@@ -36,8 +37,8 @@ import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
 public class SoftDeleteHandlerV1 extends DeleteHandlerV1 {
 
     @Inject
-    public SoftDeleteHandlerV1(AtlasTypeRegistry typeRegistry) {
-        super(typeRegistry, false, true);
+    public SoftDeleteHandlerV1(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        super(graph, typeRegistry, false, true);
     }
 
     @Override

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityGraphDiscoveryV2.java

@@ -24,6 +24,7 @@ import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasStruct;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
@@ -56,11 +57,13 @@ import static org.apache.atlas.repository.store.graph.v2.EntityGraphMapper.valid
 public class AtlasEntityGraphDiscoveryV2 implements EntityGraphDiscovery {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityGraphDiscoveryV2.class);
 
+    private final AtlasGraph                  graph;
     private final AtlasTypeRegistry           typeRegistry;
     private final EntityGraphDiscoveryContext discoveryContext;
     private final EntityGraphMapper           entityGraphMapper;
 
-    public AtlasEntityGraphDiscoveryV2(AtlasTypeRegistry typeRegistry, EntityStream entityStream, EntityGraphMapper entityGraphMapper) {
+    public AtlasEntityGraphDiscoveryV2(AtlasGraph graph, AtlasTypeRegistry typeRegistry, EntityStream entityStream, EntityGraphMapper entityGraphMapper) {
+        this.graph             = graph;
         this.typeRegistry      = typeRegistry;
         this.discoveryContext  = new EntityGraphDiscoveryContext(typeRegistry, entityStream);
         this.entityGraphMapper = entityGraphMapper;
@@ -189,8 +192,8 @@ public class AtlasEntityGraphDiscoveryV2 implements EntityGraphDiscovery {
     protected void resolveReferences() throws AtlasBaseException {
         MetricRecorder metric = RequestContext.get().startMetricRecord("resolveReferences");
 
-        EntityResolver[] entityResolvers = new EntityResolver[] { new IDBasedEntityResolver(typeRegistry, entityGraphMapper),
-                                                                  new UniqAttrBasedEntityResolver(typeRegistry, entityGraphMapper)
+        EntityResolver[] entityResolvers = new EntityResolver[] { new IDBasedEntityResolver(this.graph, typeRegistry),
+                                                                  new UniqAttrBasedEntityResolver(this.graph, typeRegistry, entityGraphMapper)
                                                                 };
 
         for (EntityResolver resolver : entityResolvers) {

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2.java

@@ -43,6 +43,7 @@ import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.AtlasBaseTypeDef;
 import org.apache.atlas.repository.Constants;
 import org.apache.atlas.repository.graph.GraphHelper;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscovery;
@@ -100,22 +101,22 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV2.class);
     private static final Logger PERF_LOG = AtlasPerfTracer.getPerfLogger("store.EntityStore");
 
-
+    private final AtlasGraph                graph;
     private final DeleteHandlerDelegate     deleteDelegate;
     private final AtlasTypeRegistry         typeRegistry;
     private final IAtlasEntityChangeNotifier entityChangeNotifier;
     private final EntityGraphMapper         entityGraphMapper;
     private final EntityGraphRetriever      entityRetriever;
 
-
     @Inject
-    public AtlasEntityStoreV2(DeleteHandlerDelegate deleteDelegate, AtlasTypeRegistry typeRegistry,
+    public AtlasEntityStoreV2(AtlasGraph graph, DeleteHandlerDelegate deleteDelegate, AtlasTypeRegistry typeRegistry,
                               IAtlasEntityChangeNotifier entityChangeNotifier, EntityGraphMapper entityGraphMapper) {
+        this.graph                = graph;
         this.deleteDelegate       = deleteDelegate;
         this.typeRegistry         = typeRegistry;
         this.entityChangeNotifier = entityChangeNotifier;
         this.entityGraphMapper    = entityGraphMapper;
-        this.entityRetriever      = new EntityGraphRetriever(typeRegistry);
+        this.entityRetriever      = new EntityGraphRetriever(graph, typeRegistry);
     }
 
     @Override
@@ -129,7 +130,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME);
         }
 
-        List<String> ret = AtlasGraphUtilsV2.findEntityGUIDsByType(typename);
+        List<String> ret = AtlasGraphUtilsV2.findEntityGUIDsByType(graph, typename);
 
         if (LOG.isDebugEnabled()) {
             LOG.debug("<== getEntityGUIDS({})", typename);
@@ -151,7 +152,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getById({}, {})", guid, isMinExtInfo);
         }
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry, ignoreRelationships);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);
 
         AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(guid, isMinExtInfo);
 
@@ -175,7 +176,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getHeaderById({})", guid);
         }
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
 
         AtlasEntityHeader ret = entityRetriever.toAtlasEntityHeader(guid);
 
@@ -205,7 +206,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getByIds({}, {})", guids, isMinExtInfo);
         }
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry, ignoreRelationships);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);
 
         AtlasEntitiesWithExtInfo ret = entityRetriever.toAtlasEntitiesWithExtInfo(guids, isMinExtInfo);
 
@@ -231,7 +232,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getEntitiesByUniqueAttributes({}, {})", entityType.getTypeName(), uniqueAttributes);
         }
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry, ignoreRelationships);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);
 
         AtlasEntitiesWithExtInfo ret = entityRetriever.getEntitiesByUniqueAttributes(entityType.getTypeName(), uniqueAttributes, isMinExtInfo);
 
@@ -262,9 +263,9 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getByUniqueAttribute({}, {})", entityType.getTypeName(), uniqAttributes);
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(entityType, uniqAttributes);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(graph, entityType, uniqAttributes);
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry, ignoreRelationships);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry, ignoreRelationships);
 
         AtlasEntityWithExtInfo ret = entityRetriever.toAtlasEntityWithExtInfo(entityVertex, isMinExtInfo);
 
@@ -289,9 +290,9 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> getEntityHeaderByUniqueAttributes({}, {})", entityType.getTypeName(), uniqAttributes);
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(entityType, uniqAttributes);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(graph, entityType, uniqAttributes);
 
-        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(typeRegistry);
+        EntityGraphRetriever entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
 
         AtlasEntityHeader ret = entityRetriever.toAtlasEntityHeader(entityVertex);
 
@@ -322,7 +323,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             LOG.debug("==> checkState({})", request);
         }
 
-        EntityStateChecker entityStateChecker = new EntityStateChecker(typeRegistry);
+        EntityStateChecker entityStateChecker = new EntityStateChecker(graph, typeRegistry);
 
         AtlasCheckStateResult ret = entityStateChecker.checkState(request);
 
@@ -372,7 +373,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
                 throw new AtlasBaseException(AtlasErrorCode.UNKNOWN_TYPENAME, objectId.getTypeName());
             }
 
-            guid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(typeRegistry.getEntityTypeByName(objectId.getTypeName()), objectId.getUniqueAttributes());
+            guid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(graph, typeRegistry.getEntityTypeByName(objectId.getTypeName()), objectId.getUniqueAttributes());
         }
 
         AtlasEntity entity = updatedEntityInfo.getEntity();
@@ -394,7 +395,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "no entity to update.");
         }
 
-        String      guid   = AtlasGraphUtilsV2.getGuidByUniqueAttributes(entityType, uniqAttributes);
+        String      guid   = AtlasGraphUtilsV2.getGuidByUniqueAttributes(graph, entityType, uniqAttributes);
         AtlasEntity entity = updatedEntityInfo.getEntity();
 
         entity.setGuid(guid);
@@ -463,7 +464,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         }
 
         Collection<AtlasVertex> deletionCandidates = new ArrayList<>();
-        AtlasVertex             vertex             = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex             vertex             = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (vertex != null) {
             AtlasEntityHeader entityHeader = entityRetriever.toAtlasEntityHeaderWithClassifications(vertex);
@@ -497,7 +498,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         Collection<AtlasVertex> deletionCandidates = new ArrayList<>();
 
         for (String guid : guids) {
-            AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
             if (vertex == null) {
                 if (LOG.isDebugEnabled()) {
@@ -539,7 +540,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         Collection<AtlasVertex> purgeCandidates = new ArrayList<>();
 
         for (String guid : guids) {
-            AtlasVertex vertex = AtlasGraphUtilsV2.findDeletedByGuid(guid);
+            AtlasVertex vertex = AtlasGraphUtilsV2.findDeletedByGuid(graph, guid);
 
             if (vertex == null) {
                 // Entity does not exist - treat as non-error, since the caller
@@ -572,7 +573,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         }
 
         Collection<AtlasVertex> deletionCandidates = new ArrayList<>();
-        AtlasVertex             vertex             = AtlasGraphUtilsV2.findByUniqueAttributes(entityType, uniqAttributes);
+        AtlasVertex             vertex             = AtlasGraphUtilsV2.findByUniqueAttributes(graph, entityType, uniqAttributes);
 
         if (vertex != null) {
             AtlasEntityHeader entityHeader = entityRetriever.toAtlasEntityHeaderWithClassifications(vertex);
@@ -599,7 +600,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
     @Override
     @GraphTransaction
     public String getGuidByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> uniqAttributes) throws AtlasBaseException{
-        return AtlasGraphUtilsV2.getGuidByUniqueAttributes(entityType, uniqAttributes);
+        return AtlasGraphUtilsV2.getGuidByUniqueAttributes(graph, entityType, uniqAttributes);
     }
 
     @Override
@@ -619,7 +620,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
 
         GraphTransactionInterceptor.lockObjectAndReleasePostCommit(guid);
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -670,7 +671,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
 
         GraphTransactionInterceptor.lockObjectAndReleasePostCommit(guid);
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -715,7 +716,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
         GraphTransactionInterceptor.lockObjectAndReleasePostCommit(guids);
 
         for (String guid : guids) {
-            AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
             if (entityVertex == null) {
                 throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -840,7 +841,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
     @Override
     @GraphTransaction
     public String setClassifications(AtlasEntityHeaders entityHeaders) {
-        ClassificationAssociator.Updater associator = new ClassificationAssociator.Updater(typeRegistry, this);
+        ClassificationAssociator.Updater associator = new ClassificationAssociator.Updater(graph, typeRegistry, this);
         return associator.setClassifications(entityHeaders.getGuidHeaderMap());
     }
 
@@ -860,7 +861,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "businessAttributes is null/empty");
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -923,7 +924,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "businessAttributes is null/empty");
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -958,7 +959,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "guid is null/empty");
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -1021,7 +1022,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "labels is null/empty");
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -1060,7 +1061,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_PARAMETERS, "labels is null/empty");
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -1247,7 +1248,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
     private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, EntityGraphMapper entityGraphMapper, boolean isPartialUpdate) throws AtlasBaseException {
         MetricRecorder metric = RequestContext.get().startMetricRecord("preCreateOrUpdate");
 
-        EntityGraphDiscovery        graphDiscoverer  = new AtlasEntityGraphDiscoveryV2(typeRegistry, entityStream, entityGraphMapper);
+        EntityGraphDiscovery        graphDiscoverer  = new AtlasEntityGraphDiscoveryV2(graph, typeRegistry, entityStream, entityGraphMapper);
         EntityGraphDiscoveryContext discoveryContext = graphDiscoverer.discoverEntities();
         EntityMutationContext       context          = new EntityMutationContext(discoveryContext);
         RequestContext              requestContext   = RequestContext.get();
@@ -1428,7 +1429,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
      */
     private void validateEntityAssociations(String guid, List<AtlasClassification> classifications) throws AtlasBaseException {
         List<String>    entityClassifications = getClassificationNames(guid);
-        String          entityTypeName        = AtlasGraphUtilsV2.getTypeNameFromGuid(guid);
+        String          entityTypeName        = AtlasGraphUtilsV2.getTypeNameFromGuid(graph, guid);
         AtlasEntityType entityType            = typeRegistry.getEntityTypeByName(entityTypeName);
 
         for (AtlasClassification classification : classifications) {
@@ -1607,7 +1608,7 @@ public class AtlasEntityStoreV2 implements AtlasEntityStore {
             String vertexKey = typeName + "_" + uniqueAttrName + "_" + uniqueAttrValue;
             AtlasVertex atlasVertex = vertexCache.get(vertexKey);
             if (atlasVertex == null) {
-                atlasVertex = AtlasGraphUtilsV2.findByTypeAndUniquePropertyName(typeName, uniqueAttrName, uniqueAttrValue);
+                atlasVertex = AtlasGraphUtilsV2.findByTypeAndUniquePropertyName(graph, typeName, uniqueAttrName, uniqueAttrValue);
             }
 
             if (atlasVertex == null) {

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasGraphUtilsV2.java

@@ -45,6 +45,7 @@ import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasStructType.AtlasAttribute;
 import org.apache.atlas.type.AtlasType;
 import org.apache.atlas.util.FileUtils;
+import org.apache.atlas.utils.AtlasPerfMetrics;
 import org.apache.atlas.utils.AtlasPerfMetrics.MetricRecorder;
 import org.apache.commons.collections.MapUtils;
 import org.apache.commons.configuration.Configuration;
@@ -268,9 +269,12 @@ public class AtlasGraphUtilsV2 {
 
         return returnType.cast(property);
     }
-
     public static AtlasVertex getVertexByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> attrValues) throws AtlasBaseException {
-        AtlasVertex vertex = findByUniqueAttributes(entityType, attrValues);
+        return getVertexByUniqueAttributes(getGraphInstance(), entityType, attrValues);
+    }
+
+    public static AtlasVertex getVertexByUniqueAttributes(AtlasGraph graph, AtlasEntityType entityType, Map<String, Object> attrValues) throws AtlasBaseException {
+        AtlasVertex vertex = findByUniqueAttributes(graph, entityType, attrValues);
 
         if (vertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, entityType.getTypeName(),
@@ -281,11 +285,20 @@ public class AtlasGraphUtilsV2 {
     }
 
     public static String getGuidByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> attrValues) throws AtlasBaseException {
-        AtlasVertex vertexByUniqueAttributes = getVertexByUniqueAttributes(entityType, attrValues);
+        AtlasVertex vertexByUniqueAttributes = getVertexByUniqueAttributes(getGraphInstance(), entityType, attrValues);
+        return getIdFromVertex(vertexByUniqueAttributes);
+    }
+
+    public static String getGuidByUniqueAttributes(AtlasGraph graph, AtlasEntityType entityType, Map<String, Object> attrValues) throws AtlasBaseException {
+        AtlasVertex vertexByUniqueAttributes = getVertexByUniqueAttributes(graph, entityType, attrValues);
         return getIdFromVertex(vertexByUniqueAttributes);
     }
 
     public static AtlasVertex findByUniqueAttributes(AtlasEntityType entityType, Map<String, Object> attrValues) {
+        return findByUniqueAttributes(getGraphInstance(), entityType, attrValues);
+    }
+
+    public static AtlasVertex findByUniqueAttributes(AtlasGraph graph, AtlasEntityType entityType, Map<String, Object> attrValues) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findByUniqueAttributes");
 
         AtlasVertex vertex = null;
@@ -300,22 +313,22 @@ public class AtlasGraphUtilsV2 {
                     continue;
                 }
 
-                if (canUseIndexQuery(entityType, attribute.getName())) {
-                    vertex = AtlasGraphUtilsV2.getAtlasVertexFromIndexQuery(entityType, attribute, attrValue);
+                if (canUseIndexQuery(graph, entityType, attribute.getName())) {
+                    vertex = AtlasGraphUtilsV2.getAtlasVertexFromIndexQuery(graph, entityType, attribute, attrValue);
                 } else {
                     if (USE_UNIQUE_INDEX_PROPERTY_TO_FIND_ENTITY && attribute.getVertexUniquePropertyName() != null) {
-                        vertex = AtlasGraphUtilsV2.findByTypeAndUniquePropertyName(entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
+                        vertex = AtlasGraphUtilsV2.findByTypeAndUniquePropertyName(graph, entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
 
                         // if no instance of given typeName is found, try to find an instance of type's sub-type
                         if (vertex == null && !entityType.getAllSubTypes().isEmpty()) {
-                            vertex = AtlasGraphUtilsV2.findBySuperTypeAndUniquePropertyName(entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
+                            vertex = AtlasGraphUtilsV2.findBySuperTypeAndUniquePropertyName(graph, entityType.getTypeName(), attribute.getVertexUniquePropertyName(), attrValue);
                         }
                     } else {
-                        vertex = AtlasGraphUtilsV2.findByTypeAndPropertyName(entityType.getTypeName(), attribute.getVertexPropertyName(), attrValue);
+                        vertex = AtlasGraphUtilsV2.findByTypeAndPropertyName(graph, entityType.getTypeName(), attribute.getVertexPropertyName(), attrValue);
 
                         // if no instance of given typeName is found, try to find an instance of type's sub-type
                         if (vertex == null && !entityType.getAllSubTypes().isEmpty()) {
-                            vertex = AtlasGraphUtilsV2.findBySuperTypeAndPropertyName(entityType.getTypeName(), attribute.getVertexPropertyName(), attrValue);
+                            vertex = AtlasGraphUtilsV2.findBySuperTypeAndPropertyName(graph, entityType.getTypeName(), attribute.getVertexPropertyName(), attrValue);
                         }
                     }
 
@@ -338,11 +351,16 @@ public class AtlasGraphUtilsV2 {
     }
 
     public static AtlasVertex findByGuid(String guid) {
+        return findByGuid(getGraphInstance(), guid);
+    }
+
+    public static AtlasVertex findByGuid(AtlasGraph graph, String guid) {
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("findByGuid");
+
         AtlasVertex ret = GraphTransactionInterceptor.getVertexFromCache(guid);
 
         if (ret == null) {
-            AtlasGraphQuery query = getGraphInstance().query()
-                    .has(Constants.GUID_PROPERTY_KEY, guid);
+            AtlasGraphQuery query = graph.query().has(Constants.GUID_PROPERTY_KEY, guid);
 
             Iterator<AtlasVertex> results = query.vertices().iterator();
 
@@ -353,14 +371,15 @@ public class AtlasGraphUtilsV2 {
             }
         }
 
+        RequestContext.get().endMetricRecord(metric);
         return ret;
     }
 
-    public static AtlasVertex findDeletedByGuid(String guid) {
+    public static AtlasVertex findDeletedByGuid(AtlasGraph graph, String guid) {
         AtlasVertex ret = GraphTransactionInterceptor.getVertexFromCache(guid);
 
         if (ret == null) {
-            AtlasGraphQuery query = getGraphInstance().query()
+            AtlasGraphQuery query = graph.query()
                     .has(Constants.GUID_PROPERTY_KEY, guid)
                     .has(STATE_PROPERTY_KEY, Status.DELETED.name());
 
@@ -376,20 +395,23 @@ public class AtlasGraphUtilsV2 {
         return ret;
     }
 
-    public static String getTypeNameFromGuid(String guid) {
+    public static String getTypeNameFromGuid(AtlasGraph graph, String guid) {
         String ret = null;
 
         if (StringUtils.isNotEmpty(guid)) {
-            AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex vertex = AtlasGraphUtilsV2.findByGuid(graph, guid);
 
             ret = (vertex != null) ? AtlasGraphUtilsV2.getTypeName(vertex) : null;
         }
 
         return ret;
     }
-
     public static boolean typeHasInstanceVertex(String typeName) throws AtlasBaseException {
-        AtlasGraphQuery query = getGraphInstance()
+        return typeHasInstanceVertex(getGraphInstance(), typeName);
+    }
+
+    public static boolean typeHasInstanceVertex(AtlasGraph graph, String typeName) throws AtlasBaseException {
+        AtlasGraphQuery query = graph
                 .query()
                 .has(TYPE_NAME_PROPERTY_KEY, AtlasGraphQuery.ComparisionOperator.EQUAL, typeName);
 
@@ -405,9 +427,13 @@ public class AtlasGraphUtilsV2 {
     }
 
     public static AtlasVertex findByTypeAndUniquePropertyName(String typeName, String propertyName, Object attrVal) {
+        return findByTypeAndUniquePropertyName(getGraphInstance(), typeName, propertyName, attrVal);
+    }
+
+    public static AtlasVertex findByTypeAndUniquePropertyName(AtlasGraph graph, String typeName, String propertyName, Object attrVal) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findByTypeAndUniquePropertyName");
 
-        AtlasGraphQuery query = getGraphInstance().query()
+        AtlasGraphQuery query = graph.query()
                                                     .has(ENTITY_TYPE_PROPERTY_KEY, typeName)
                                                     .has(propertyName, attrVal);
 
@@ -420,10 +446,10 @@ public class AtlasGraphUtilsV2 {
         return vertex;
     }
 
-    public static AtlasVertex findBySuperTypeAndUniquePropertyName(String typeName, String propertyName, Object attrVal) {
+    public static AtlasVertex findBySuperTypeAndUniquePropertyName(AtlasGraph graph, String typeName, String propertyName, Object attrVal) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findBySuperTypeAndUniquePropertyName");
 
-        AtlasGraphQuery query = getGraphInstance().query()
+        AtlasGraphQuery query = graph.query()
                                                     .has(Constants.SUPER_TYPES_PROPERTY_KEY, typeName)
                                                     .has(propertyName, attrVal);
 
@@ -436,10 +462,10 @@ public class AtlasGraphUtilsV2 {
         return vertex;
     }
 
-    public static AtlasVertex findByTypeAndPropertyName(String typeName, String propertyName, Object attrVal) {
+    public static AtlasVertex findByTypeAndPropertyName(AtlasGraph graph, String typeName, String propertyName, Object attrVal) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findByTypeAndPropertyName");
 
-        AtlasGraphQuery query = getGraphInstance().query()
+        AtlasGraphQuery query = graph.query()
                                                     .has(ENTITY_TYPE_PROPERTY_KEY, typeName)
                                                     .has(propertyName, attrVal)
                                                     .has(STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
@@ -453,10 +479,10 @@ public class AtlasGraphUtilsV2 {
         return vertex;
     }
 
-    public static AtlasVertex findBySuperTypeAndPropertyName(String typeName, String propertyName, Object attrVal) {
+    public static AtlasVertex findBySuperTypeAndPropertyName(AtlasGraph graph, String typeName, String propertyName, Object attrVal) {
         MetricRecorder metric = RequestContext.get().startMetricRecord("findBySuperTypeAndPropertyName");
 
-        AtlasGraphQuery query = getGraphInstance().query()
+        AtlasGraphQuery query = graph.query()
                                                     .has(Constants.SUPER_TYPES_PROPERTY_KEY, typeName)
                                                     .has(propertyName, attrVal)
                                                     .has(STATE_PROPERTY_KEY, AtlasEntity.Status.ACTIVE.name());
@@ -470,14 +496,12 @@ public class AtlasGraphUtilsV2 {
         return vertex;
     }
 
-    public int getOpenTransactions() {
-        Set openTransactions = getGraphInstance().getOpenTransactions();
-
-        return (openTransactions != null) ? openTransactions.size() : 0;
+    public static List<String> findEntityGUIDsByType(String typename, SortOrder sortOrder) {
+        return findEntityGUIDsByType(getGraphInstance(), typename, sortOrder);
     }
 
-    public static List<String> findEntityGUIDsByType(String typename, SortOrder sortOrder) {
-        AtlasGraphQuery query = getGraphInstance().query()
+    public static List<String> findEntityGUIDsByType(AtlasGraph graph, String typename, SortOrder sortOrder) {
+        AtlasGraphQuery query = graph.query()
                                                   .has(ENTITY_TYPE_PROPERTY_KEY, typename);
         if (sortOrder != null) {
             AtlasGraphQuery.SortOrder qrySortOrder = sortOrder == SortOrder.ASCENDING ? ASC : DESC;
@@ -498,8 +522,8 @@ public class AtlasGraphUtilsV2 {
         return ret;
     }
 
-    public static Iterator<AtlasVertex> findActiveEntityVerticesByType(String typename) {
-        return findActiveEntityVerticesByType(getGraphInstance(), typename);
+    public static List<String> findEntityGUIDsByType(AtlasGraph graph, String typename) {
+        return findEntityGUIDsByType(graph, typename, null);
     }
 
     public static Iterator<AtlasVertex> findActiveEntityVerticesByType(AtlasGraph graph, String typename) {
@@ -510,12 +534,12 @@ public class AtlasGraphUtilsV2 {
         return query.vertices().iterator();
     }
 
-    public static List<String> findEntityGUIDsByType(String typename) {
-        return findEntityGUIDsByType(typename, null);
+    public static boolean relationshipTypeHasInstanceEdges(String typeName) throws AtlasBaseException {
+        return relationshipTypeHasInstanceEdges(getGraphInstance(), typeName);
     }
 
-    public static boolean relationshipTypeHasInstanceEdges(String typeName) throws AtlasBaseException {
-        AtlasGraphQuery query = getGraphInstance()
+    public static boolean relationshipTypeHasInstanceEdges(AtlasGraph graph, String typeName) throws AtlasBaseException {
+        AtlasGraphQuery query = graph
                 .query()
                 .has(TYPE_NAME_PROPERTY_KEY, AtlasGraphQuery.ComparisionOperator.EQUAL, typeName);
 
@@ -584,7 +608,7 @@ public class AtlasGraphUtilsV2 {
         return element.getProperty(STATE_PROPERTY_KEY, String.class);
     }
 
-    private static boolean canUseIndexQuery(AtlasEntityType entityType, String attributeName) {
+    private static boolean canUseIndexQuery(AtlasGraph graph, AtlasEntityType entityType, String attributeName) {
         boolean ret = false;
 
         if (USE_INDEX_QUERY_TO_FIND_ENTITY_BY_UNIQUE_ATTRIBUTES) {
@@ -593,7 +617,7 @@ public class AtlasGraphUtilsV2 {
             ret = typeAndSubTypesQryStr.length() <= SearchProcessor.MAX_QUERY_STR_LENGTH_TYPES;
 
             if (ret) {
-                Set<String> indexSet = getGraphInstance().getVertexIndexKeys();
+                Set<String> indexSet = graph.getVertexIndexKeys();
                 try {
                     ret = indexSet.contains(entityType.getQualifiedAttributeName(attributeName));
                 }
@@ -606,9 +630,9 @@ public class AtlasGraphUtilsV2 {
         return ret;
     }
 
-    private static AtlasVertex getAtlasVertexFromIndexQuery(AtlasEntityType entityType, AtlasAttribute attribute, Object attrVal) {
+    private static AtlasVertex getAtlasVertexFromIndexQuery(AtlasGraph graph, AtlasEntityType entityType, AtlasAttribute attribute, Object attrVal) {
         String          propertyName = attribute.getVertexPropertyName();
-        AtlasIndexQuery query        = getIndexQuery(entityType, propertyName, attrVal.toString());
+        AtlasIndexQuery query        = getIndexQuery(graph, entityType, propertyName, attrVal.toString());
 
         for (Iterator<Result> iter = query.vertices(); iter.hasNext(); ) {
             Result      result = iter.next();
@@ -645,7 +669,7 @@ public class AtlasGraphUtilsV2 {
         return null;
     }
 
-    private static AtlasIndexQuery getIndexQuery(AtlasEntityType entityType, String propertyName, String value) {
+    private static AtlasIndexQuery getIndexQuery(AtlasGraph graph, AtlasEntityType entityType, String propertyName, String value) {
         StringBuilder sb = new StringBuilder();
 
         sb.append(INDEX_SEARCH_PREFIX + "\"").append(TYPE_NAME_PROPERTY_KEY).append("\":").append(entityType.getTypeAndAllSubTypesQryStr())
@@ -654,7 +678,7 @@ public class AtlasGraphUtilsV2 {
                 .append(" AND ")
                 .append(INDEX_SEARCH_PREFIX + "\"").append(STATE_PROPERTY_KEY).append("\":ACTIVE");
 
-        return getGraphInstance().indexQuery(Constants.VERTEX_INDEX, sb.toString());
+        return graph.indexQuery(Constants.VERTEX_INDEX, sb.toString());
     }
 
     public static String getIndexSearchPrefix() {

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2.java

@@ -40,6 +40,7 @@ import org.apache.atlas.repository.RepositoryException;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
 import org.apache.atlas.repository.store.graph.v1.DeleteHandlerDelegate;
@@ -91,20 +92,23 @@ import static org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2.getTy
 @Component
 public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasRelationshipStoreV2.class);
-
     private static final Long DEFAULT_RELATIONSHIP_VERSION = 0L;
+
+    private final AtlasGraph graph;
     private boolean notificationsEnabled = NOTIFICATION_RELATIONSHIPS_ENABLED.getBoolean();
 
     private final AtlasTypeRegistry         typeRegistry;
     private final EntityGraphRetriever      entityRetriever;
     private final DeleteHandlerDelegate     deleteDelegate;
-    private final GraphHelper               graphHelper = GraphHelper.getInstance();
+    private final GraphHelper               graphHelper;
     private final IAtlasEntityChangeNotifier entityChangeNotifier;
 
     @Inject
-    public AtlasRelationshipStoreV2(AtlasTypeRegistry typeRegistry, DeleteHandlerDelegate deleteDelegate, IAtlasEntityChangeNotifier entityChangeNotifier) {
+    public AtlasRelationshipStoreV2(AtlasGraph graph, AtlasTypeRegistry typeRegistry, DeleteHandlerDelegate deleteDelegate, IAtlasEntityChangeNotifier entityChangeNotifier) {
+        this.graph                = graph;
         this.typeRegistry         = typeRegistry;
-        this.entityRetriever      = new EntityGraphRetriever(typeRegistry);
+        this.graphHelper          = new GraphHelper(graph);
+        this.entityRetriever      = new EntityGraphRetriever(graph, typeRegistry);
         this.deleteDelegate       = deleteDelegate;
         this.entityChangeNotifier = entityChangeNotifier;
     }
@@ -274,7 +278,7 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
             throw new AtlasBaseException(AtlasErrorCode.RELATIONSHIP_ALREADY_DELETED, guid);
         }
 
-        String            relationShipType = GraphHelper.getTypeName(edge);
+        String            relationShipType = graphHelper.getTypeName(edge);
         AtlasEntityHeader end1Entity       = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getOutVertex());
         AtlasEntityHeader end2Entity       = entityRetriever.toAtlasEntityHeaderWithClassifications(edge.getInVertex());
 
@@ -721,7 +725,7 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
             String              guid             = end.getGuid();
             String              typeName         = end.getTypeName();
             Map<String, Object> uniqueAttributes = end.getUniqueAttributes();
-            AtlasVertex         endVertex        = AtlasGraphUtilsV2.findByGuid(guid);
+            AtlasVertex         endVertex        = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
             if (!AtlasTypeUtil.isValidGuid(guid) || endVertex == null) {
                 throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -729,7 +733,7 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
             } else if (MapUtils.isNotEmpty(uniqueAttributes)) {
                 AtlasEntityType entityType = typeRegistry.getEntityTypeByName(typeName);
 
-                if (AtlasGraphUtilsV2.findByUniqueAttributes(entityType, uniqueAttributes) == null) {
+                if (AtlasGraphUtilsV2.findByUniqueAttributes(this.graph, entityType, uniqueAttributes) == null) {
                     throw new AtlasBaseException(AtlasErrorCode.INSTANCE_BY_UNIQUE_ATTRIBUTE_NOT_FOUND, typeName, uniqueAttributes.toString());
                 }
             } else {
@@ -801,11 +805,11 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
         AtlasVertex ret = null;
 
         if (StringUtils.isNotEmpty(endPoint.getGuid())) {
-            ret = AtlasGraphUtilsV2.findByGuid(endPoint.getGuid());
+            ret = AtlasGraphUtilsV2.findByGuid(this.graph, endPoint.getGuid());
         } else if (StringUtils.isNotEmpty(endPoint.getTypeName()) && MapUtils.isNotEmpty(endPoint.getUniqueAttributes())) {
             AtlasEntityType entityType = typeRegistry.getEntityTypeByName(endPoint.getTypeName());
 
-            ret = AtlasGraphUtilsV2.findByUniqueAttributes(entityType, endPoint.getUniqueAttributes());
+            ret = AtlasGraphUtilsV2.findByUniqueAttributes(this.graph, entityType, endPoint.getUniqueAttributes());
         }
 
         return ret;
@@ -888,7 +892,7 @@ public class AtlasRelationshipStoreV2 implements AtlasRelationshipStore {
         String typeName = objectId.getTypeName();
 
         if (StringUtils.isBlank(typeName)) {
-            typeName = AtlasGraphUtilsV2.getTypeNameFromGuid(objectId.getGuid());
+            typeName = AtlasGraphUtilsV2.getTypeNameFromGuid(this.graph, objectId.getGuid());
         }
 
         return typeName;

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/BulkImporterImpl.java

@@ -26,6 +26,7 @@ import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.BulkImporter;
@@ -51,10 +52,12 @@ public class BulkImporterImpl implements BulkImporter {
     private static final Logger LOG = LoggerFactory.getLogger(AtlasEntityStoreV2.class);
 
     private final AtlasEntityStore entityStore;
+    private final AtlasGraph atlasGraph;
     private AtlasTypeRegistry typeRegistry;
 
     @Inject
-    public BulkImporterImpl(AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
+    public BulkImporterImpl(AtlasGraph atlasGraph, AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
+        this.atlasGraph = atlasGraph;
         this.entityStore = entityStore;
         this.typeRegistry = typeRegistry;
     }
@@ -65,9 +68,9 @@ public class BulkImporterImpl implements BulkImporter {
 
         if (importResult.getRequest().getOptions() != null &&
                 importResult.getRequest().getOptions().containsKey(AtlasImportRequest.OPTION_KEY_MIGRATION)) {
-            importStrategy = new MigrationImport(new AtlasGraphProvider(), this.typeRegistry);
+            importStrategy = new MigrationImport(this.atlasGraph, new AtlasGraphProvider(), this.typeRegistry);
         } else {
-            importStrategy = new RegularImport(this.entityStore, this.typeRegistry);
+            importStrategy = new RegularImport(this.atlasGraph, this.entityStore, this.typeRegistry);
         }
 
         LOG.info("BulkImportImpl: {}", importStrategy.getClass().getSimpleName());
@@ -121,14 +124,14 @@ public class BulkImporterImpl implements BulkImporter {
         }
     }
 
-    public static void updateVertexGuid(AtlasTypeRegistry typeRegistry, EntityGraphRetriever entityGraphRetriever, AtlasEntity entity) {
+    public static void updateVertexGuid(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, EntityGraphRetriever entityGraphRetriever, AtlasEntity entity) {
         String entityGuid = entity.getGuid();
         AtlasObjectId objectId = entityGraphRetriever.toAtlasObjectIdWithoutGuid(entity);
 
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
         String vertexGuid = null;
         try {
-            vertexGuid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(entityType, objectId.getUniqueAttributes());
+            vertexGuid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(atlasGraph, entityType, objectId.getUniqueAttributes());
         } catch (AtlasBaseException e) {
             LOG.warn("Entity: {}: Does not exist!", objectId);
             return;
@@ -138,7 +141,7 @@ public class BulkImporterImpl implements BulkImporter {
             return;
         }
 
-        AtlasVertex v = AtlasGraphUtilsV2.findByGuid(vertexGuid);
+        AtlasVertex v = AtlasGraphUtilsV2.findByGuid(atlasGraph, vertexGuid);
         if (v == null) {
             return;
         }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociator.java

@@ -23,6 +23,8 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasEntityHeaders;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
+import org.apache.atlas.repository.graph.AtlasGraphProvider;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.type.AtlasEntityType;
@@ -49,11 +51,15 @@ public class ClassificationAssociator {
         private final EntityAuditRepository auditRepository;
         private final EntityGraphRetriever entityRetriever;
 
-        public Retriever(AtlasTypeRegistry typeRegistry, EntityAuditRepository auditRepository) {
-            this.entityRetriever = new EntityGraphRetriever(typeRegistry);
+        public Retriever(AtlasGraph graph, AtlasTypeRegistry typeRegistry, EntityAuditRepository auditRepository) {
+            this.entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
             this.auditRepository = auditRepository;
         }
 
+        public Retriever(AtlasTypeRegistry typeRegistry, EntityAuditRepository auditRepository) {
+            this(AtlasGraphProvider.getGraphInstance(), typeRegistry, auditRepository);
+        }
+
         Retriever(EntityGraphRetriever entityGraphRetriever, EntityAuditRepository auditRepository) {
             this.entityRetriever = entityGraphRetriever;
             this.auditRepository = auditRepository;
@@ -104,16 +110,21 @@ public class ClassificationAssociator {
         static final String PROCESS_DELETE = "Delete";
         static final String JSONIFY_STRING_FORMAT = "\"%s\",";
 
+        private final AtlasGraph graph;
         private final AtlasTypeRegistry typeRegistry;
         private final AtlasEntityStore entitiesStore;
         private final EntityGraphRetriever entityRetriever;
-        private StringBuilder actionSummary = new StringBuilder();
-        private Map<String, String> typeNameUniqueAttributeNameMap = new HashMap<>();
+        private final StringBuilder actionSummary = new StringBuilder();
 
-        public Updater(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
+        public Updater(AtlasGraph graph, AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
+            this.graph = graph;
             this.typeRegistry = typeRegistry;
             this.entitiesStore = entitiesStore;
-            entityRetriever = new EntityGraphRetriever(typeRegistry);
+            entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
+        }
+
+        public Updater(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
+            this(AtlasGraphProvider.getGraphInstance(), typeRegistry, entitiesStore);
         }
 
         public String setClassifications(Map<String, AtlasEntityHeader> map) {
@@ -246,7 +257,7 @@ public class ClassificationAssociator {
 
         AtlasEntityHeader getByUniqueAttributes(AtlasEntityType entityType, String qualifiedName, Map<String, Object> attrValues) {
             try {
-                AtlasVertex vertex = AtlasGraphUtilsV2.findByUniqueAttributes(entityType, attrValues);
+                AtlasVertex vertex = AtlasGraphUtilsV2.findByUniqueAttributes(this.graph, entityType, attrValues);
                 if (vertex == null) {
                     return null;
                 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphMapper.java

@@ -92,7 +92,6 @@ import static org.apache.atlas.repository.Constants.*;
 import static org.apache.atlas.repository.graph.GraphHelper.getCollectionElementsUsingRelationship;
 import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdge;
 import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertex;
-import static org.apache.atlas.repository.graph.GraphHelper.getDefaultRemovePropagations;
 import static org.apache.atlas.repository.graph.GraphHelper.getDelimitedClassificationNames;
 import static org.apache.atlas.repository.graph.GraphHelper.getLabels;
 import static org.apache.atlas.repository.graph.GraphHelper.getMapElementsProperty;
@@ -130,7 +129,7 @@ public class EntityGraphMapper {
     private static final boolean ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES = AtlasConfiguration.ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES.getBoolean();
     private static final boolean CLASSIFICATION_PROPAGATION_DEFAULT                  = AtlasConfiguration.CLASSIFICATION_PROPAGATION_DEFAULT.getBoolean();
 
-    private final GraphHelper               graphHelper = GraphHelper.getInstance();
+    private final GraphHelper               graphHelper;
     private final AtlasGraph                graph;
     private final DeleteHandlerDelegate     deleteDelegate;
     private final AtlasTypeRegistry         typeRegistry;
@@ -141,16 +140,17 @@ public class EntityGraphMapper {
     private final IFullTextMapper fullTextMapperV2;
 
     @Inject
-    public EntityGraphMapper(DeleteHandlerDelegate deleteDelegate, AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph,
+    public EntityGraphMapper(DeleteHandlerDelegate deleteDelegate, AtlasTypeRegistry typeRegistry, AtlasGraph graph,
                              AtlasRelationshipStore relationshipStore, IAtlasEntityChangeNotifier entityChangeNotifier,
                              AtlasInstanceConverter instanceConverter, IFullTextMapper fullTextMapperV2) {
+        this.graphHelper          = new GraphHelper(graph);
         this.deleteDelegate       = deleteDelegate;
         this.typeRegistry         = typeRegistry;
-        this.graph                = atlasGraph;
+        this.graph                = graph;
         this.relationshipStore    = relationshipStore;
         this.entityChangeNotifier = entityChangeNotifier;
         this.instanceConverter    = instanceConverter;
-        this.entityRetriever      = new EntityGraphRetriever(typeRegistry);
+        this.entityRetriever      = new EntityGraphRetriever(graph, typeRegistry);
         this.fullTextMapperV2     = fullTextMapperV2;
     }
 
@@ -361,12 +361,12 @@ public class EntityGraphMapper {
 
         updateLabels(vertex, labels);
 
-        entityChangeNotifier.onLabelsUpdatedFromEntity(GraphHelper.getGuid(vertex), addedLabels, removedLabels);
+        entityChangeNotifier.onLabelsUpdatedFromEntity(graphHelper.getGuid(vertex), addedLabels, removedLabels);
     }
 
     public void addLabels(AtlasVertex vertex, Set<String> labels) throws AtlasBaseException {
         if (CollectionUtils.isNotEmpty(labels)) {
-            final Set<String> existingLabels = GraphHelper.getLabels(vertex);
+            final Set<String> existingLabels = graphHelper.getLabels(vertex);
             final Set<String> updatedLabels;
 
             if (CollectionUtils.isEmpty(existingLabels)) {
@@ -378,14 +378,14 @@ public class EntityGraphMapper {
             if (!updatedLabels.equals(existingLabels)) {
                 updateLabels(vertex, updatedLabels);
                 updatedLabels.removeAll(existingLabels);
-                entityChangeNotifier.onLabelsUpdatedFromEntity(GraphHelper.getGuid(vertex), updatedLabels, null);
+                entityChangeNotifier.onLabelsUpdatedFromEntity(graphHelper.getGuid(vertex), updatedLabels, null);
             }
         }
     }
 
     public void removeLabels(AtlasVertex vertex, Set<String> labels) throws AtlasBaseException {
         if (CollectionUtils.isNotEmpty(labels)) {
-            final Set<String> existingLabels = GraphHelper.getLabels(vertex);
+            final Set<String> existingLabels = graphHelper.getLabels(vertex);
             Set<String> updatedLabels;
 
             if (CollectionUtils.isNotEmpty(existingLabels)) {
@@ -395,7 +395,7 @@ public class EntityGraphMapper {
                 if (!updatedLabels.equals(existingLabels)) {
                     updateLabels(vertex, updatedLabels);
                     existingLabels.removeAll(updatedLabels);
-                    entityChangeNotifier.onLabelsUpdatedFromEntity(GraphHelper.getGuid(vertex), null, existingLabels);
+                    entityChangeNotifier.onLabelsUpdatedFromEntity(graphHelper.getGuid(vertex), null, existingLabels);
                 }
             }
         }
@@ -695,9 +695,10 @@ public class EntityGraphMapper {
     }
 
     private void mapAttribute(AtlasAttribute attribute, Object attrValue, AtlasVertex vertex, EntityOperation op, EntityMutationContext context) throws AtlasBaseException {
+        boolean isDeletedEntity = context.isDeletedEntity(vertex);
+        AtlasType         attrType     = attribute.getAttributeType();
         if (attrValue == null) {
             AtlasAttributeDef attributeDef = attribute.getAttributeDef();
-            AtlasType         attrType     = attribute.getAttributeType();
 
             if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE) {
                 if (attributeDef.getDefaultValue() != null) {
@@ -712,10 +713,13 @@ public class EntityGraphMapper {
             }
         }
 
+        if (attrType.getTypeCategory() == TypeCategory.PRIMITIVE || attrType.getTypeCategory() == TypeCategory.ENUM) {
+            mapPrimitiveValue(vertex, attribute, attrValue, isDeletedEntity);
+        } else {
             AttributeMutationContext ctx = new AttributeMutationContext(op, vertex, attribute, attrValue);
-
             mapToVertexByTypeCategory(ctx, context);
         }
+    }
 
     private Object mapToVertexByTypeCategory(AttributeMutationContext ctx, EntityMutationContext context) throws AtlasBaseException {
         if (ctx.getOp() == CREATE && ctx.getValue() == null) {
@@ -906,7 +910,7 @@ public class EntityGraphMapper {
         if (inverseUpdated) {
             RequestContext requestContext = RequestContext.get();
 
-            if (!requestContext.isDeletedEntity(GraphHelper.getGuid(inverseVertex))) {
+            if (!requestContext.isDeletedEntity(graphHelper.getGuid(inverseVertex))) {
                 updateModificationMetadata(inverseVertex);
 
                 requestContext.recordEntityUpdate(entityRetriever.toAtlasEntityHeader(inverseVertex));
@@ -960,7 +964,7 @@ public class EntityGraphMapper {
             return;
         }
 
-        String parentGuid = GraphHelper.getGuid(inverseVertex);
+        String parentGuid = graphHelper.getGuid(inverseVertex);
         if(StringUtils.isEmpty(parentGuid)) {
             return;
         }
@@ -997,9 +1001,13 @@ public class EntityGraphMapper {
     }
 
     private Object mapPrimitiveValue(AttributeMutationContext ctx, EntityMutationContext context) {
-        boolean isIndexableStrAttr = ctx.getAttributeDef().getIsIndexable() && ctx.getAttrType() instanceof AtlasBuiltInTypes.AtlasStringType;
+        return mapPrimitiveValue(ctx.getReferringVertex(), ctx.getAttribute(), ctx.getValue(), context.isDeletedEntity(ctx.referringVertex));
+    }
+
+    private Object mapPrimitiveValue(AtlasVertex vertex, AtlasAttribute attribute, Object valueFromEntity, boolean isDeletedEntity) {
+        boolean isIndexableStrAttr = attribute.getAttributeDef().getIsIndexable() && attribute.getAttributeType() instanceof AtlasBuiltInTypes.AtlasStringType;
 
-        Object ret = ctx.getValue();
+        Object ret = valueFromEntity;
 
         // Janus bug, when an indexed string attribute has a value longer than a certain length then the reverse indexed key generated by JanusGraph
         // exceeds the HBase row length's hard limit (Short.MAX). This trimming and hashing procedure is to circumvent that limitation
@@ -1024,26 +1032,26 @@ public class EntityGraphMapper {
                 }
 
                 if (trimmedLength < value.length()) {
-                    LOG.warn("Length of indexed attribute {} is {} characters, longer than safe-limit {}; trimming to {} - attempt #{}", ctx.getAttribute().getQualifiedName(), value.length(), INDEXED_STR_SAFE_LEN, trimmedLength, requestContext.getAttemptCount());
+                    LOG.warn("Length of indexed attribute {} is {} characters, longer than safe-limit {}; trimming to {} - attempt #{}", attribute.getQualifiedName(), value.length(), INDEXED_STR_SAFE_LEN, trimmedLength, requestContext.getAttemptCount());
 
                     String checksumSuffix = ":" + DigestUtils.shaHex(value); // Storing SHA checksum in case verification is needed after retrieval
 
                     ret = value.substring(0, trimmedLength - checksumSuffix.length()) + checksumSuffix;
                 } else {
-                    LOG.warn("Length of indexed attribute {} is {} characters, longer than safe-limit {}", ctx.getAttribute().getQualifiedName(), value.length(), INDEXED_STR_SAFE_LEN);
+                    LOG.warn("Length of indexed attribute {} is {} characters, longer than safe-limit {}", attribute.getQualifiedName(), value.length(), INDEXED_STR_SAFE_LEN);
                 }
             }
         }
 
-        AtlasGraphUtilsV2.setEncodedProperty(ctx.getReferringVertex(), ctx.getVertexProperty(), ret);
+        AtlasGraphUtilsV2.setEncodedProperty(vertex, attribute.getVertexPropertyName(), ret);
 
-        String uniqPropName = ctx.getAttribute() != null ? ctx.getAttribute().getVertexUniquePropertyName() : null;
+        String uniqPropName = attribute != null ? attribute.getVertexUniquePropertyName() : null;
 
         if (uniqPropName != null) {
-            if (context.isDeletedEntity(ctx.getReferringVertex()) || AtlasGraphUtilsV2.getState(ctx.getReferringVertex()) == DELETED) {
-                ctx.getReferringVertex().removeProperty(uniqPropName);
+            if (isDeletedEntity || AtlasGraphUtilsV2.getState(vertex) == DELETED) {
+                vertex.removeProperty(uniqPropName);
             } else {
-                AtlasGraphUtilsV2.setEncodedProperty(ctx.getReferringVertex(), uniqPropName, ret);
+                AtlasGraphUtilsV2.setEncodedProperty(vertex, uniqPropName, ret);
             }
         }
 
@@ -1203,7 +1211,7 @@ public class EntityGraphMapper {
 
                     ret = getOrCreateRelationship(fromVertex, toVertex, relationshipName, relationshipAttributes);
 
-                    boolean isCreated = GraphHelper.getCreatedTime(ret) == RequestContext.get().getRequestTime();
+                    boolean isCreated = graphHelper.getCreatedTime(ret) == RequestContext.get().getRequestTime();
 
                     if (isCreated) {
                         // if relationship did not exist before and new relationship was created
@@ -1549,7 +1557,7 @@ public class EntityGraphMapper {
         return null;
     }
 
-    private static void setAssignedGuid(Object val, EntityMutationContext context) {
+    private void setAssignedGuid(Object val, EntityMutationContext context) {
         if (val != null) {
             Map<String, String> guidAssignements = context.getGuidAssignments();
 
@@ -1566,7 +1574,7 @@ public class EntityGraphMapper {
                     AtlasVertex vertex = context.getDiscoveryContext().getResolvedEntityVertex(objId);
 
                     if (vertex != null) {
-                        assignedGuid = GraphHelper.getGuid(vertex);
+                        assignedGuid = graphHelper.getGuid(vertex);
                     }
                 }
 
@@ -1589,7 +1597,7 @@ public class EntityGraphMapper {
                     AtlasVertex vertex = context.getDiscoveryContext().getResolvedEntityVertex(new AtlasObjectId(mapObjId));
 
                     if (vertex != null) {
-                        assignedGuid = GraphHelper.getGuid(vertex);
+                        assignedGuid = graphHelper.getGuid(vertex);
                     }
                 }
 
@@ -1907,7 +1915,7 @@ public class EntityGraphMapper {
                 }
 
                 if (removePropagations == null) {
-                    removePropagations = getDefaultRemovePropagations();
+                    removePropagations = graphHelper.getDefaultRemovePropagations();
 
                     classification.setRemovePropagationsOnEntityDelete(removePropagations);
                 }
@@ -2007,7 +2015,7 @@ public class EntityGraphMapper {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_CLASSIFICATION_PARAMS, "delete", entityGuid);
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(entityGuid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(this.graph, entityGuid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, entityGuid);
@@ -2021,7 +2029,7 @@ public class EntityGraphMapper {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_CLASSIFICATION_PARAMS, "delete", entityGuid);
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(entityGuid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(this.graph, entityGuid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, entityGuid);
@@ -2092,7 +2100,7 @@ public class EntityGraphMapper {
     }
 
     private AtlasEntity updateClassificationText(AtlasVertex vertex) throws AtlasBaseException {
-        String guid        = GraphHelper.getGuid(vertex);
+        String guid        = graphHelper.getGuid(vertex);
         AtlasEntity entity = instanceConverter.getAndCacheEntity(guid, ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
 
         vertex.setProperty(CLASSIFICATION_TEXT_KEY, fullTextMapperV2.getClassificationTextForEntity(entity));
@@ -2105,7 +2113,7 @@ public class EntityGraphMapper {
             return;
         }
 
-        String guid = GraphHelper.getGuid(vertex);
+        String guid = graphHelper.getGuid(vertex);
         AtlasEntity entity = instanceConverter.getAndCacheEntity(guid, ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
         List<String> classificationNames = new ArrayList<>();
         List<String> propagatedClassificationNames = new ArrayList<>();
@@ -2163,7 +2171,7 @@ public class EntityGraphMapper {
             throw new AtlasBaseException(AtlasErrorCode.INVALID_CLASSIFICATION_PARAMS, "update", guid);
         }
 
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -2315,7 +2323,7 @@ public class EntityGraphMapper {
         }
 
         for (AtlasVertex vertex : notificationVertices) {
-            String      entityGuid = GraphHelper.getGuid(vertex);
+            String      entityGuid = graphHelper.getGuid(vertex);
             AtlasEntity entity     = instanceConverter.getAndCacheEntity(entityGuid, ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
 
             if (isActive(entity)) {
@@ -2369,7 +2377,7 @@ public class EntityGraphMapper {
     }
 
     public void deleteClassifications(String guid) throws AtlasBaseException {
-        AtlasVertex instanceVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex instanceVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
         if (instanceVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -2408,7 +2416,7 @@ public class EntityGraphMapper {
         if (vertex != null) {
             RequestContext req = RequestContext.get();
 
-            if (!req.isUpdatedEntity(GraphHelper.getGuid(vertex))) {
+            if (!req.isUpdatedEntity(graphHelper.getGuid(vertex))) {
                 updateModificationMetadata(vertex);
 
                 req.recordEntityUpdate(entityRetriever.toAtlasEntityHeader(vertex));
@@ -2535,7 +2543,7 @@ public class EntityGraphMapper {
                     continue;
                 }
 
-                if (value.length() > CUSTOM_ATTRIBUTE_VALUE_MAX_LENGTH) {
+                if (!key.startsWith(CUSTOM_ATTRIBUTE_KEY_SPECIAL_PREFIX) && value.length() > CUSTOM_ATTRIBUTE_VALUE_MAX_LENGTH) {
                     throw new AtlasBaseException(AtlasErrorCode.INVALID_CUSTOM_ATTRIBUTE_VALUE, value, String.valueOf(CUSTOM_ATTRIBUTE_VALUE_MAX_LENGTH));
                 }
             }
@@ -2563,7 +2571,7 @@ public class EntityGraphMapper {
 
         if(CollectionUtils.isNotEmpty(propagatedVertices)) {
             for(AtlasVertex vertex : propagatedVertices) {
-                AtlasEntity entity = instanceConverter.getAndCacheEntity(GraphHelper.getGuid(vertex), ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
+                AtlasEntity entity = instanceConverter.getAndCacheEntity(graphHelper.getGuid(vertex), ENTITY_CHANGE_NOTIFY_IGNORE_RELATIONSHIP_ATTRIBUTES);
 
                 if (isActive(entity)) {
                     vertex.setProperty(CLASSIFICATION_TEXT_KEY, fullTextMapperV2.getClassificationTextForEntity(entity));

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityGraphRetriever.java

@@ -43,6 +43,7 @@ import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasElement;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasArrayType;
 import org.apache.atlas.type.AtlasBuiltInTypes.AtlasObjectIdType;
@@ -131,18 +132,21 @@ public class EntityGraphRetriever {
     public static final String QUALIFIED_NAME = "qualifiedName";
 
     private static final TypeReference<List<TimeBoundary>> TIME_BOUNDARIES_LIST_TYPE = new TypeReference<List<TimeBoundary>>() {};
-    private static final GraphHelper graphHelper = GraphHelper.getInstance();
+    private final GraphHelper graphHelper;
 
     private final AtlasTypeRegistry typeRegistry;
 
     private final boolean ignoreRelationshipAttr;
+    private final AtlasGraph graph;
 
     @Inject
-    public EntityGraphRetriever(AtlasTypeRegistry typeRegistry) {
-        this(typeRegistry, false);
+    public EntityGraphRetriever(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        this(graph, typeRegistry, false);
     }
 
-    public EntityGraphRetriever(AtlasTypeRegistry typeRegistry, boolean ignoreRelationshipAttr) {
+    public EntityGraphRetriever(AtlasGraph graph, AtlasTypeRegistry typeRegistry, boolean ignoreRelationshipAttr) {
+        this.graph                  = graph;
+        this.graphHelper            = new GraphHelper(graph);
         this.typeRegistry           = typeRegistry;
         this.ignoreRelationshipAttr = ignoreRelationshipAttr;
     }
@@ -365,7 +369,7 @@ public class EntityGraphRetriever {
     }
 
     public AtlasVertex getEntityVertex(String guid) throws AtlasBaseException {
-        AtlasVertex ret = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex ret = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
         if (ret == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);
@@ -419,7 +423,7 @@ public class EntityGraphRetriever {
         if (entityType != null) {
             for (Map<String, Object> uniqAttributes : uniqueAttributesList) {
                 try {
-                    AtlasVertex vertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(entityType, uniqAttributes);
+                    AtlasVertex vertex = AtlasGraphUtilsV2.getVertexByUniqueAttributes(this.graph, entityType, uniqAttributes);
 
                     if (vertex != null) {
                         AtlasEntity entity = mapVertexToAtlasEntity(vertex, ret, isMinExtInfo);
@@ -442,7 +446,7 @@ public class EntityGraphRetriever {
     public void evaluateClassificationPropagation(AtlasVertex classificationVertex, List<AtlasVertex> entitiesToAddPropagation, List<AtlasVertex> entitiesToRemovePropagation) {
         if (classificationVertex != null) {
             String            entityGuid         = getClassificationEntityGuid(classificationVertex);
-            AtlasVertex       entityVertex       = AtlasGraphUtilsV2.findByGuid(entityGuid);
+            AtlasVertex       entityVertex       = AtlasGraphUtilsV2.findByGuid(this.graph, entityGuid);
             String            classificationId   = classificationVertex.getIdForDisplay();
             List<AtlasVertex> propagatedEntities = getAllPropagatedEntityVertices(classificationVertex);
             List<AtlasVertex> impactedEntities   = getImpactedVerticesV2(entityVertex, null, classificationId);
@@ -472,7 +476,7 @@ public class EntityGraphRetriever {
             for (AtlasVertex classificationVertex : classificationVertices) {
                 String            classificationId      = classificationVertex.getIdForDisplay();
                 String            sourceEntityId        = getClassificationEntityGuid(classificationVertex);
-                AtlasVertex       sourceEntityVertex    = AtlasGraphUtilsV2.findByGuid(sourceEntityId);
+                AtlasVertex       sourceEntityVertex    = AtlasGraphUtilsV2.findByGuid(this.graph, sourceEntityId);
                 List<AtlasVertex> entitiesPropagatingTo = getImpactedVerticesV2(sourceEntityVertex, relationshipGuidToExclude, classificationId);
 
                 ret.put(classificationVertex, entitiesPropagatingTo);
@@ -582,12 +586,12 @@ public class EntityGraphRetriever {
         }
 
         if (AtlasTypeUtil.isAssignedGuid(objId)) {
-            ret = AtlasGraphUtilsV2.findByGuid(objId.getGuid());
+            ret = AtlasGraphUtilsV2.findByGuid(this.graph, objId.getGuid());
         } else {
             AtlasEntityType     entityType     = typeRegistry.getEntityTypeByName(objId.getTypeName());
             Map<String, Object> uniqAttributes = objId.getUniqueAttributes();
 
-            ret = AtlasGraphUtilsV2.getVertexByUniqueAttributes(entityType, uniqAttributes);
+            ret = AtlasGraphUtilsV2.getVertexByUniqueAttributes(this.graph, entityType, uniqAttributes);
         }
 
         if (ret == null) {
@@ -828,10 +832,10 @@ public class EntityGraphRetriever {
         return ret;
     }
 
-    public List<AtlasTermAssignmentHeader> mapAssignedTerms(AtlasVertex entityVertex) throws AtlasBaseException {
+    public List<AtlasTermAssignmentHeader> mapAssignedTerms(AtlasVertex entityVertex) {
         List<AtlasTermAssignmentHeader> ret = new ArrayList<>();
 
-        Iterable edges = entityVertex.query().direction(AtlasEdgeDirection.IN).label(TERM_ASSIGNMENT_LABEL).edges();
+        Iterable edges = entityVertex.query().direction(AtlasEdgeDirection.IN).label(Constants.INTERNAL_PROPERTY_KEY_PREFIX + TERM_ASSIGNMENT_LABEL).edges();
 
         if (edges != null) {
             for (final AtlasEdge edge : (Iterable<AtlasEdge>) edges) {

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/EntityStateChecker.java

@@ -28,13 +28,13 @@ import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphHelper;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.commons.collections.CollectionUtils;
-import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.stereotype.Component;
@@ -55,13 +55,15 @@ import static org.apache.atlas.repository.graph.GraphHelper.getDelimitedClassifi
 public final class EntityStateChecker {
     private static final Logger LOG = LoggerFactory.getLogger(EntityStateChecker.class);
 
+    private final AtlasGraph           graph;
     private final AtlasTypeRegistry    typeRegistry;
     private final EntityGraphRetriever entityRetriever;
 
     @Inject
-    public EntityStateChecker(AtlasTypeRegistry typeRegistry) {
+    public EntityStateChecker(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        this.graph           = graph;
         this.typeRegistry    = typeRegistry;
-        this.entityRetriever = new EntityGraphRetriever(typeRegistry);
+        this.entityRetriever = new EntityGraphRetriever(graph, typeRegistry);
     }
 
 
@@ -140,7 +142,7 @@ public final class EntityStateChecker {
      * @throws AtlasBaseException
      */
     public AtlasEntityState checkEntityState(String guid, boolean fixIssues, AtlasCheckStateResult result) throws AtlasBaseException {
-        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(guid);
+        AtlasVertex entityVertex = AtlasGraphUtilsV2.findByGuid(this.graph, guid);
 
         if (entityVertex == null) {
             throw new AtlasBaseException(AtlasErrorCode.INSTANCE_GUID_NOT_FOUND, guid);

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/IDBasedEntityResolver.java

@@ -22,6 +22,7 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasEntity;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
@@ -35,12 +36,12 @@ import org.slf4j.LoggerFactory;
 public class IDBasedEntityResolver implements EntityResolver {
     private static final Logger LOG = LoggerFactory.getLogger(IDBasedEntityResolver.class);
 
+    private final AtlasGraph        graph;
     private final AtlasTypeRegistry typeRegistry;
-    private final EntityGraphMapper entityGraphMapper;
 
-    public IDBasedEntityResolver(AtlasTypeRegistry typeRegistry, EntityGraphMapper entityGraphMapper) {
+    public IDBasedEntityResolver(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        this.graph             = graph;
         this.typeRegistry      = typeRegistry;
-        this.entityGraphMapper = entityGraphMapper;
     }
 
     public EntityGraphDiscoveryContext resolveEntityReferences(EntityGraphDiscoveryContext context) throws AtlasBaseException {
@@ -52,7 +53,7 @@ public class IDBasedEntityResolver implements EntityResolver {
 
         for (String guid : context.getReferencedGuids()) {
             boolean isAssignedGuid = AtlasTypeUtil.isAssignedGuid(guid);
-            AtlasVertex vertex = isAssignedGuid ? AtlasGraphUtilsV2.findByGuid(guid) : null;
+            AtlasVertex vertex = isAssignedGuid ? AtlasGraphUtilsV2.findByGuid(this.graph, guid) : null;
 
             if (vertex == null && !RequestContext.get().isImportInProgress()) { // if not found in the store, look if the entity is present in the stream
                 AtlasEntity entity = entityStream.getByGuid(guid);
@@ -64,7 +65,7 @@ public class IDBasedEntityResolver implements EntityResolver {
                         throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), entity.getTypeName());
                     }
 
-                    vertex = AtlasGraphUtilsV2.findByUniqueAttributes(entityType, entity.getAttributes());
+                    vertex = AtlasGraphUtilsV2.findByUniqueAttributes(this.graph, entityType, entity.getAttributes());
                 } else if (!isAssignedGuid) { // for local-guids, entity must be in the stream
                     throw new AtlasBaseException(AtlasErrorCode.REFERENCED_ENTITY_NOT_FOUND, guid);
                 }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/UniqAttrBasedEntityResolver.java

@@ -22,6 +22,7 @@ import org.apache.atlas.RequestContext;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
 import org.apache.atlas.model.instance.AtlasObjectId;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.EntityGraphDiscoveryContext;
 import org.apache.atlas.repository.store.graph.EntityResolver;
@@ -36,10 +37,12 @@ import java.util.List;
 public class UniqAttrBasedEntityResolver implements EntityResolver {
     private static final Logger LOG = LoggerFactory.getLogger(UniqAttrBasedEntityResolver.class);
 
+    private final AtlasGraph        graph;
     private final AtlasTypeRegistry typeRegistry;
     private final EntityGraphMapper entityGraphMapper;
 
-    public UniqAttrBasedEntityResolver(AtlasTypeRegistry typeRegistry, EntityGraphMapper entityGraphMapper) {
+    public UniqAttrBasedEntityResolver(AtlasGraph graph, AtlasTypeRegistry typeRegistry, EntityGraphMapper entityGraphMapper) {
+        this.graph             = graph;
         this.typeRegistry      = typeRegistry;
         this.entityGraphMapper = entityGraphMapper;
     }
@@ -61,7 +64,7 @@ public class UniqAttrBasedEntityResolver implements EntityResolver {
                 throw new AtlasBaseException(AtlasErrorCode.TYPE_NAME_INVALID, TypeCategory.ENTITY.name(), objId.getTypeName());
             }
 
-            AtlasVertex vertex = AtlasGraphUtilsV2.findByUniqueAttributes(entityType, objId.getUniqueAttributes());
+            AtlasVertex vertex = AtlasGraphUtilsV2.findByUniqueAttributes(this.graph, entityType, objId.getUniqueAttributes());
 
             if (vertex == null && RequestContext.get().isCreateShellEntityForNonExistingReference()) {
                 vertex = entityGraphMapper.createShellEntityVertex(objId, context);

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/bulkimport/MigrationImport.java

@@ -28,7 +28,6 @@ import org.apache.atlas.repository.converters.AtlasInstanceConverter;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.migration.DataMigrationStatusService;
-import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
 import org.apache.atlas.repository.store.graph.v1.DeleteHandlerDelegate;
 import org.apache.atlas.repository.store.graph.v2.AtlasEntityStoreV2;
@@ -46,15 +45,14 @@ import org.slf4j.LoggerFactory;
 public class MigrationImport extends ImportStrategy {
     private static final Logger LOG = LoggerFactory.getLogger(MigrationImport.class);
 
+    private final AtlasGraph graph;
+    private final AtlasGraphProvider graphProvider;
     private final AtlasTypeRegistry typeRegistry;
-    private AtlasGraph atlasGraph;
-    private EntityGraphRetriever entityGraphRetriever;
-    private EntityGraphMapper entityGraphMapper;
-    private AtlasEntityStore entityStore;
 
-    public MigrationImport(AtlasGraphProvider atlasGraphProvider, AtlasTypeRegistry typeRegistry) {
+    public MigrationImport(AtlasGraph graph, AtlasGraphProvider graphProvider, AtlasTypeRegistry typeRegistry) {
+        this.graph = graph;
+        this.graphProvider = graphProvider;
         this.typeRegistry = typeRegistry;
-        setupEntityStore(atlasGraphProvider, typeRegistry);
         LOG.info("MigrationImport: Using bulkLoading...");
     }
 
@@ -72,7 +70,7 @@ public class MigrationImport extends ImportStrategy {
         long index = 0;
         int streamSize = entityStream.size();
         EntityMutationResponse ret = new EntityMutationResponse();
-        EntityCreationManager creationManager = createEntityCreationManager(atlasGraph, importResult, dataMigrationStatusService);
+        EntityCreationManager creationManager = createEntityCreationManager(importResult, dataMigrationStatusService);
 
         try {
             LOG.info("Migration Import: Size: {}: Starting...", streamSize);
@@ -95,16 +93,24 @@ public class MigrationImport extends ImportStrategy {
         return dataMigrationStatusService;
     }
 
-    private EntityCreationManager createEntityCreationManager(AtlasGraph threadedAtlasGraph,
-                                                              AtlasImportResult importResult,
+    private EntityCreationManager createEntityCreationManager(AtlasImportResult importResult,
                                                               DataMigrationStatusService dataMigrationStatusService) {
-        atlasGraph = threadedAtlasGraph;
+        AtlasGraph graphBulk = graphProvider.getBulkLoading();
+
+        EntityGraphRetriever entityGraphRetriever = new EntityGraphRetriever(this.graph, typeRegistry);
+        EntityGraphRetriever entityGraphRetrieverBulk = new EntityGraphRetriever(graphBulk, typeRegistry);
+
+        AtlasEntityStoreV2 entityStore = createEntityStore(this.graph, typeRegistry);
+        AtlasEntityStoreV2 entityStoreBulk = createEntityStore(graphBulk, typeRegistry);
+
         int batchSize = importResult.getRequest().getOptionKeyBatchSize();
         int numWorkers = getNumWorkers(importResult.getRequest().getOptionKeyNumWorkers());
 
         EntityConsumerBuilder consumerBuilder =
-                new EntityConsumerBuilder(threadedAtlasGraph, entityStore, entityGraphRetriever, typeRegistry, batchSize);
+                new EntityConsumerBuilder(typeRegistry, this.graph, entityStore, entityGraphRetriever, graphBulk,
+                        entityStoreBulk, entityGraphRetrieverBulk, batchSize);
 
+        LOG.info("MigrationImport: EntityCreationManager: Created!");
         return new EntityCreationManager(consumerBuilder, batchSize, numWorkers, importResult, dataMigrationStatusService);
     }
 
@@ -114,17 +120,17 @@ public class MigrationImport extends ImportStrategy {
         return ret;
     }
 
-    private void setupEntityStore(AtlasGraphProvider atlasGraphProvider, AtlasTypeRegistry typeRegistry) {
-        this.entityGraphRetriever = new EntityGraphRetriever(typeRegistry);
-        this.atlasGraph = atlasGraphProvider.getBulkLoading();
-        DeleteHandlerDelegate deleteDelegate = new DeleteHandlerDelegate(typeRegistry);
-
+    private AtlasEntityStoreV2 createEntityStore(AtlasGraph graph, AtlasTypeRegistry typeRegistry) {
+        FullTextMapperV2Nop fullTextMapperV2 = new FullTextMapperV2Nop();
         IAtlasEntityChangeNotifier entityChangeNotifier = new EntityChangeNotifierNop();
-        AtlasRelationshipStore relationshipStore = new AtlasRelationshipStoreV2(typeRegistry, deleteDelegate, entityChangeNotifier);
+        DeleteHandlerDelegate deleteDelegate = new DeleteHandlerDelegate(graph, typeRegistry);
         AtlasFormatConverters formatConverters = new AtlasFormatConverters(typeRegistry);
-        AtlasInstanceConverter instanceConverter = new AtlasInstanceConverter(typeRegistry, formatConverters);
-        this.entityGraphMapper = new EntityGraphMapper(deleteDelegate, typeRegistry, atlasGraph, relationshipStore, entityChangeNotifier, instanceConverter, new FullTextMapperV2Nop());
-        this.entityStore = new AtlasEntityStoreV2(deleteDelegate, typeRegistry, entityChangeNotifier, entityGraphMapper);
+
+        AtlasInstanceConverter instanceConverter = new AtlasInstanceConverter(graph, typeRegistry, formatConverters);
+        AtlasRelationshipStore relationshipStore = new AtlasRelationshipStoreV2(graph, typeRegistry, deleteDelegate, entityChangeNotifier);
+        EntityGraphMapper entityGraphMapper = new EntityGraphMapper(deleteDelegate, typeRegistry, graph, relationshipStore, entityChangeNotifier, instanceConverter, fullTextMapperV2);
+
+        return new AtlasEntityStoreV2(graph, deleteDelegate, typeRegistry, entityChangeNotifier, entityGraphMapper);
     }
 
     private void shutdownEntityCreationManager(EntityCreationManager creationManager) {

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/bulkimport/RegularImport.java

@@ -31,6 +31,7 @@ import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.repository.Constants;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.graphdb.AtlasSchemaViolationException;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
@@ -57,15 +58,18 @@ import static org.apache.atlas.repository.store.graph.v2.BulkImporterImpl.update
 public class RegularImport extends ImportStrategy {
     private static final Logger LOG = LoggerFactory.getLogger(RegularImport.class);
     private static final int MAX_ATTEMPTS = 3;
+
+    private final AtlasGraph graph;
     private final AtlasEntityStore entityStore;
     private final AtlasTypeRegistry typeRegistry;
     private final EntityGraphRetriever entityGraphRetriever;
     private boolean directoryBasedImportConfigured;
 
-    public RegularImport(AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
+    public RegularImport(AtlasGraph graph, AtlasEntityStore entityStore, AtlasTypeRegistry typeRegistry) {
+        this.graph       = graph;
         this.entityStore = entityStore;
         this.typeRegistry = typeRegistry;
-        this.entityGraphRetriever = new EntityGraphRetriever(typeRegistry);
+        this.entityGraphRetriever = new EntityGraphRetriever(graph, typeRegistry);
         this.directoryBasedImportConfigured = StringUtils.isNotEmpty(AtlasConfiguration.IMPORT_TEMP_DIRECTORY.getString());
     }
 
@@ -156,7 +160,7 @@ public class RegularImport extends ImportStrategy {
         AtlasEntityType entityType = typeRegistry.getEntityTypeByName(entity.getTypeName());
         String vertexGuid = null;
         try {
-            vertexGuid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(entityType, objectId.getUniqueAttributes());
+            vertexGuid = AtlasGraphUtilsV2.getGuidByUniqueAttributes(this.graph, entityType, objectId.getUniqueAttributes());
         } catch (AtlasBaseException e) {
             LOG.warn("Entity: {}: Does not exist!", objectId);
             return;
@@ -166,7 +170,7 @@ public class RegularImport extends ImportStrategy {
             return;
         }
 
-        AtlasVertex v = AtlasGraphUtilsV2.findByGuid(vertexGuid);
+        AtlasVertex v = AtlasGraphUtilsV2.findByGuid(this.graph, vertexGuid);
         if (v == null) {
             return;
         }

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/bulkimport/pc/EntityConsumer.java

@@ -30,7 +30,9 @@ import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.v2.AtlasEntityStreamForImport;
 import org.apache.atlas.repository.store.graph.v2.BulkImporterImpl;
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
+import org.apache.atlas.repository.store.graph.v2.EntityStream;
 import org.apache.atlas.type.AtlasTypeRegistry;
+import org.apache.atlas.utils.AtlasPerfMetrics;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.collections.MapUtils;
 import org.slf4j.Logger;
@@ -49,23 +51,30 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
     private AtomicLong counter = new AtomicLong(1);
     private AtomicLong currentBatch = new AtomicLong(1);
 
-    private final AtlasGraph atlasGraph;
-    private final AtlasEntityStore entityStoreV2;
+    private AtlasGraph atlasGraph;
+    private final AtlasEntityStore entityStore;
+    private final AtlasGraph atlasGraphBulk;
+    private final AtlasEntityStore entityStoreBulk;
     private final AtlasTypeRegistry typeRegistry;
-    private final EntityGraphRetriever entityGraphRetriever;
+    private final EntityGraphRetriever entityRetrieverBulk;
 
     private List<AtlasEntity.AtlasEntityWithExtInfo> entityBuffer = new ArrayList<>();
     private List<EntityMutationResponse> localResults = new ArrayList<>();
 
-    public EntityConsumer(AtlasGraph atlasGraph, AtlasEntityStore entityStore,
-                          EntityGraphRetriever entityGraphRetriever, AtlasTypeRegistry typeRegistry,
+    public EntityConsumer(AtlasTypeRegistry typeRegistry,
+                          AtlasGraph atlasGraph, AtlasEntityStore entityStore,
+                          AtlasGraph atlasGraphBulk, AtlasEntityStore entityStoreBulk, EntityGraphRetriever entityRetrieverBulk,
                           BlockingQueue queue, int batchSize) {
         super(queue);
+        this.typeRegistry = typeRegistry;
 
         this.atlasGraph = atlasGraph;
-        this.entityStoreV2 = entityStore;
-        this.entityGraphRetriever = entityGraphRetriever;
-        this.typeRegistry = typeRegistry;
+        this.entityStore = entityStore;
+
+        this.atlasGraphBulk = atlasGraphBulk;
+        this.entityStoreBulk = entityStoreBulk;
+        this.entityRetrieverBulk = entityRetrieverBulk;
+
         this.batchSize = batchSize;
     }
 
@@ -77,7 +86,6 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
 
         long currentCount = counter.addAndGet(delta);
         currentBatch.addAndGet(delta);
-        entityBuffer.add(entityWithExtInfo);
 
         try {
             processEntity(entityWithExtInfo, currentCount);
@@ -88,24 +96,68 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
     }
 
     private void processEntity(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo, long currentCount) {
-        try {
         RequestContext.get().setImportInProgress(true);
         RequestContext.get().setCreateShellEntityForNonExistingReference(true);
 
-            AtlasEntityStreamForImport oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);
-
+        try {
             LOG.debug("Processing: {}", currentCount);
-            EntityMutationResponse result = entityStoreV2.createOrUpdateForImportNoCommit(oneEntityStream);
-            localResults.add(result);
+            importUsingBulkEntityStore(entityWithExtInfo);
+        } catch (IllegalStateException | IllegalArgumentException e) {
+            LOG.warn("{}: {} - {}", e.getClass().getSimpleName(), entityWithExtInfo.getEntity().getTypeName(), entityWithExtInfo.getEntity().getGuid(), e);
+            importUsingRegularEntityStore(entityWithExtInfo, e);
         } catch (AtlasBaseException e) {
-            addResult(entityWithExtInfo.getEntity().getGuid());
-            LOG.warn("Exception: {}", entityWithExtInfo.getEntity().getGuid(), e);
+            LOG.warn("AtlasBaseException: {} - {}", entityWithExtInfo.getEntity().getTypeName(), entityWithExtInfo.getEntity().getGuid(), e);
         } catch (AtlasSchemaViolationException e) {
             if (LOG.isDebugEnabled()) {
                 LOG.debug("Entity: {}", entityWithExtInfo.getEntity().getGuid(), e);
             }
 
-            BulkImporterImpl.updateVertexGuid(typeRegistry, entityGraphRetriever, entityWithExtInfo.getEntity());
+            BulkImporterImpl.updateVertexGuid(this.atlasGraphBulk, typeRegistry, entityRetrieverBulk, entityWithExtInfo.getEntity());
+        }
+    }
+
+    private void importUsingBulkEntityStore(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) throws AtlasBaseException {
+        EntityStream oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);
+        EntityMutationResponse result = entityStoreBulk.createOrUpdateForImportNoCommit(oneEntityStream);
+        localResults.add(result);
+        entityBuffer.add(entityWithExtInfo);
+    }
+
+    private void importUsingRegularEntityStore(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo, Exception ex) {
+        commitValidatedEntities(ex);
+        performRegularImport(entityWithExtInfo);
+    }
+
+    private void performRegularImport(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
+        synchronized (atlasGraph) {
+            try {
+                LOG.info("Regular: EntityStore: {}: Starting...", this.counter.get());
+                AtlasEntityStreamForImport oneEntityStream = new AtlasEntityStreamForImport(entityWithExtInfo, null);
+                EntityMutationResponse result = this.entityStore.createOrUpdateForImportNoCommit(oneEntityStream);
+                atlasGraph.commit();
+                localResults.add(result);
+                dispatchResults();
+            } catch (Exception e) {
+                atlasGraph.rollback();
+                LOG.error("Regular: EntityStore: Rollback!: Entity creation using regular (non-bulk) failed! Please correct entity and re-submit!", e);
+            } finally {
+                LOG.info("Regular: EntityStore: {}: Commit: Done!", this.counter.get());
+                atlasGraph.commit();
+                addResult(entityWithExtInfo.getEntity().getGuid());
+                clear();
+                LOG.info("Regular: EntityStore: {}: Done!", this.counter.get());
+            }
+        }
+    }
+
+    private void commitValidatedEntities(Exception ex) {
+        try {
+            LOG.info("Validated Entities: Commit: Starting...");
+            rollbackPauseRetry(1, ex);
+            doCommit();
+        }
+        finally {
+            LOG.info("Validated Entities: Commit: Done!");
         }
     }
 
@@ -137,8 +189,10 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
     }
 
     private boolean commitWithRetry(int retryCount) {
+        AtlasPerfMetrics.MetricRecorder metric = RequestContext.get().startMetricRecord("commitWithRetry");
+
         try {
-            atlasGraph.commit();
+            atlasGraphBulk.commit();
             if (LOG.isDebugEnabled()) {
                 LOG.debug("Commit: Done!: Buffer: {}: Batch: {}: Counter: {}", entityBuffer.size(), currentBatch.get(), counter.get());
             }
@@ -148,14 +202,15 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
         } catch (Exception ex) {
             rollbackPauseRetry(retryCount, ex);
             return false;
+        } finally {
+            RequestContext.get().endMetricRecord(metric);
         }
     }
 
     private void rollbackPauseRetry(int retryCount, Exception ex) {
-        atlasGraph.rollback();
-        clearCache();
+        bulkGraphRollback(retryCount);
 
-        LOG.error("Rollback: Done! Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount);
+        LOG.warn("Rollback: Done! Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount);
         pause(retryCount);
         String exceptionClass = ex.getClass().getSimpleName();
         if (!exceptionClass.equals("JanusGraphException") && !exceptionClass.equals("PermanentLockingException")) {
@@ -164,14 +219,27 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
         retryProcessEntity(retryCount);
     }
 
+    private void bulkGraphRollback(int retryCount) {
+        try {
+            atlasGraphBulk.rollback();
+            clearCache();
+        } catch (Exception e) {
+            LOG.error("Rollback: Exception! Buffer: {}: Counter: {}: Retry count: {}", entityBuffer.size(), counter.get(), retryCount);
+        }
+    }
+
     private void retryProcessEntity(int retryCount) {
         if (LOG.isDebugEnabled() || retryCount > 1) {
             LOG.info("Replaying: Starting!: Buffer: {}: Retry count: {}", entityBuffer.size(), retryCount);
         }
 
-        for (AtlasEntity.AtlasEntityWithExtInfo e : entityBuffer) {
+        List<AtlasEntity.AtlasEntityWithExtInfo> localBuffer = new ArrayList<>(entityBuffer);
+        entityBuffer.clear();
+
+        for (AtlasEntity.AtlasEntityWithExtInfo e : localBuffer) {
             processEntity(e, counter.get());
         }
+
         LOG.info("Replaying: Done!: Buffer: {}: Retry count: {}", entityBuffer.size(), retryCount);
     }
 

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/bulkimport/pc/EntityConsumerBuilder.java

@@ -22,6 +22,7 @@ import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.pc.WorkItemBuilder;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
+import org.apache.atlas.repository.store.graph.v2.AtlasEntityStoreV2;
 import org.apache.atlas.repository.store.graph.v2.EntityGraphRetriever;
 import org.apache.atlas.type.AtlasTypeRegistry;
 
@@ -30,21 +31,33 @@ import java.util.concurrent.BlockingQueue;
 public class EntityConsumerBuilder implements WorkItemBuilder<EntityConsumer, AtlasEntity.AtlasEntityWithExtInfo> {
     private AtlasGraph atlasGraph;
     private AtlasEntityStore entityStore;
-    private final EntityGraphRetriever entityGraphRetriever;
+    private AtlasGraph atlasGraphBulk;
+    private AtlasEntityStore entityStoreBulk;
+    private final EntityGraphRetriever entityRetriever;
     private final AtlasTypeRegistry typeRegistry;
+    private EntityGraphRetriever entityRetrieverBulk;
     private int batchSize;
 
-    public EntityConsumerBuilder(AtlasGraph atlasGraph, AtlasEntityStore entityStore,
-                                 EntityGraphRetriever entityGraphRetriever, AtlasTypeRegistry typeRegistry, int batchSize) {
+    public EntityConsumerBuilder(AtlasTypeRegistry typeRegistry, AtlasGraph atlasGraph, AtlasEntityStoreV2 entityStore, EntityGraphRetriever entityRetriever,
+                                 AtlasGraph atlasGraphBulk, AtlasEntityStoreV2 entityStoreBulk, EntityGraphRetriever entityRetrieverBulk,
+                                 int batchSize) {
+        this.typeRegistry = typeRegistry;
+
         this.atlasGraph = atlasGraph;
         this.entityStore = entityStore;
-        this.entityGraphRetriever = entityGraphRetriever;
-        this.typeRegistry = typeRegistry;
+        this.entityRetriever = entityRetriever;
+
+        this.atlasGraphBulk = atlasGraphBulk;
+        this.entityStoreBulk = entityStoreBulk;
+        this.entityRetrieverBulk = entityRetrieverBulk;
+
         this.batchSize = batchSize;
     }
 
     @Override
     public EntityConsumer build(BlockingQueue<AtlasEntity.AtlasEntityWithExtInfo> queue) {
-        return new EntityConsumer(atlasGraph, entityStore, entityGraphRetriever, typeRegistry, queue, this.batchSize);
+        return new EntityConsumer(typeRegistry, atlasGraph, entityStore,
+                atlasGraphBulk, entityStoreBulk, entityRetrieverBulk,
+                queue, this.batchSize);
     }
 }

repository/src/test/java/org/apache/atlas/repository/impexp/ExportSkipLineageTest.java

@@ -26,6 +26,7 @@ import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.impexp.AtlasExportRequest;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.repository.AtlasTestBase;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.v1.DeleteHandlerDelegate;
 import org.apache.atlas.repository.store.graph.v2.AtlasEntityChangeNotifier;
 import org.apache.atlas.repository.store.graph.v2.AtlasEntityStoreV2;
@@ -66,6 +67,9 @@ public class ExportSkipLineageTest extends AtlasTestBase {
     @Inject
     ExportService exportService;
 
+    @Inject
+    AtlasGraph atlasGraph;
+
     private DeleteHandlerDelegate deleteDelegate = mock(DeleteHandlerDelegate.class);
     private AtlasEntityChangeNotifier mockChangeNotifier = mock(AtlasEntityChangeNotifier.class);
     private AtlasEntityStoreV2 entityStore;
@@ -76,7 +80,7 @@ public class ExportSkipLineageTest extends AtlasTestBase {
         loadHiveModel(typeDefStore, typeRegistry);
         RequestContext.get().setImportInProgress(true);
 
-        entityStore = new AtlasEntityStoreV2(deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
+        entityStore = new AtlasEntityStoreV2(atlasGraph, deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
         createEntities(entityStore, ENTITIES_SUB_DIR, new String[]{"db", "table-columns", "table-view", "table-table-lineage"});
         final String[] entityGuids = {DB_GUID, TABLE_GUID, TABLE_TABLE_GUID, TABLE_VIEW_GUID};
         verifyCreatedEntities(entityStore, entityGuids, 4);

repository/src/test/java/org/apache/atlas/repository/migration/HiveParititionTest.java

@@ -51,9 +51,9 @@ public class HiveParititionTest extends  MigrationBaseAsserts {
 
         assertTypeCountNameGuid("hive_db", 1, "parts_db", "ae30d78b-51b4-42ab-9436-8d60c8f68b95");
         assertTypeCountNameGuid("hive_process", 1, "", "");
-        assertEdges("hive_db", "parts_db", AtlasEdgeDirection.IN,1, 1, "");
-        assertEdges("hive_table", "t1", AtlasEdgeDirection.OUT, 1, 1, "hive_table_db");
-        assertEdges("hive_table", "tv1", AtlasEdgeDirection.OUT, 1, 1, "hive_table_db");
+        assertEdges("hive_db", "parts_db", AtlasEdgeDirection.IN, 1, "");
+        assertEdges("hive_table", "t1", AtlasEdgeDirection.OUT, 1, "hive_table_db");
+        assertEdges("hive_table", "tv1", AtlasEdgeDirection.OUT, 1, "hive_table_db");
 
         assertMigrationStatus(EXPECTED_TOTAL_COUNT);
     }

repository/src/test/java/org/apache/atlas/repository/migration/HiveStocksTest.java

@@ -59,9 +59,9 @@ public class HiveStocksTest extends MigrationBaseAsserts {
         assertTypeCountNameGuid("hive_storagedesc", 1, "", "294290d8-4498-4677-973c-c266d594b039");
         assertTypeCountNameGuid("Tag1", 1, "", "");
 
-        assertEdges(getVertex("hive_db", "stocks").getEdges(AtlasEdgeDirection.IN).iterator(),1, 1, "");
-        assertEdges(getVertex("hive_table", "stocks_daily").getEdges(AtlasEdgeDirection.OUT).iterator(), 1, 1, "hive_table_db");
-        assertEdges(getVertex("hive_column", "high").getEdges(AtlasEdgeDirection.OUT).iterator(), 1,1, "hive_table_columns");
+        assertEdges(getVertex("hive_db", "stocks").getEdges(AtlasEdgeDirection.IN).iterator(), 1, "");
+        assertEdges(getVertex("hive_table", "stocks_daily").getEdges(AtlasEdgeDirection.OUT).iterator(),  1, "hive_table_db");
+        assertEdges(getVertex("hive_column", "high").getEdges(AtlasEdgeDirection.OUT).iterator(), 1, "hive_table_columns");
 
         assertMigrationStatus(EXPECTED_TOTAL_COUNT);
     }

repository/src/test/java/org/apache/atlas/repository/migration/MigrationBaseAsserts.java

@@ -133,25 +133,33 @@ public class MigrationBaseAsserts {
         return iterator.hasNext() ? iterator.next() : null;
     }
 
-    protected void assertEdges(String typeName, String assetName, AtlasEdgeDirection edgeDirection, int startIdx, int expectedItems, String edgeTypeName) {
-        assertEdges(getVertex(typeName, assetName).getEdges(edgeDirection).iterator(),startIdx, expectedItems, edgeTypeName);
+    protected void assertEdges(String typeName, String assetName, AtlasEdgeDirection edgeDirection, int expectedItems, String edgeTypeName) {
+        Iterator edgeIterator = getVertex(typeName, assetName).getEdges(edgeDirection).iterator();
+        assertEdges(edgeIterator, expectedItems, edgeTypeName);
     }
 
-    protected void assertEdges(Iterator<AtlasEdge> results, int startIdx, int expectedItems, String edgeTypeName) {
+    protected void assertEdges(Iterator<AtlasEdge> results, int expectedItems, String edgeTypeNameExpected) {
         int count = 0;
         AtlasEdge e = null;
-        for (Iterator<AtlasEdge> it = results; it.hasNext() && count < startIdx; count++) {
+        boolean searchedEdgeFound = false;
+        for (Iterator<AtlasEdge> it = results; it.hasNext();) {
             e = it.next();
+            String typeName = AtlasGraphUtilsV2.getEncodedProperty(e, TYPE_NAME_PROPERTY, String.class);
+            searchedEdgeFound = StringUtils.isEmpty(edgeTypeNameExpected) || typeName.equals(edgeTypeNameExpected);
+            if (searchedEdgeFound) {
+                count++;
+                break;
+            }
         }
 
         assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, R_GUID_PROPERTY_NAME, Object.class));
         assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, "tagPropagation", Object.class));
 
-        if(StringUtils.isNotEmpty(edgeTypeName)) {
-            assertEquals(AtlasGraphUtilsV2.getEncodedProperty(e, TYPE_NAME_PROPERTY, Object.class), edgeTypeName, edgeTypeName);
+        if(StringUtils.isNotEmpty(edgeTypeNameExpected)) {
+            assertTrue(searchedEdgeFound, edgeTypeNameExpected);
         }
 
-        assertEquals(count, expectedItems, String.format("%s", edgeTypeName));
+        assertEquals(count, expectedItems, String.format("%s", edgeTypeNameExpected));
     }
 
     protected void assertEdgesWithLabel(Iterator<AtlasEdge> results, int startIdx, String edgeTypeName) {

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityStoreV2Test.java

@@ -125,7 +125,7 @@ public class AtlasEntityStoreV2Test extends AtlasEntityTestBase {
     }
     @BeforeTest
     public void init() throws Exception {
-        entityStore = new AtlasEntityStoreV2(deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
+        entityStore = new AtlasEntityStoreV2(graph, deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
         RequestContext.clear();
         RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);
 

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasEntityTestBase.java

@@ -107,7 +107,7 @@ public class AtlasEntityTestBase {
 
     @BeforeTest
     public void init() throws Exception {
-        entityStore = new AtlasEntityStoreV2(deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
+        entityStore = new AtlasEntityStoreV2(graph, deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
 
         RequestContext.clear();
         RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/AtlasRelationshipStoreV2Test.java

@@ -32,6 +32,7 @@ import org.apache.atlas.model.instance.EntityMutationResponse;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.repository.graph.AtlasGraphProvider;
 import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
@@ -86,6 +87,9 @@ public abstract class AtlasRelationshipStoreV2Test {
     @Inject
     AtlasEntityChangeNotifier entityNotifier;
 
+    @Inject
+    AtlasGraph atlasGraph;
+
     AtlasEntityStore          entityStore;
     AtlasRelationshipStore    relationshipStore;
     AtlasEntityChangeNotifier mockChangeNotifier = mock(AtlasEntityChangeNotifier.class);
@@ -124,8 +128,8 @@ public abstract class AtlasRelationshipStoreV2Test {
 
     @BeforeTest
     public void init() throws Exception {
-        entityStore       = new AtlasEntityStoreV2(deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
-        relationshipStore = new AtlasRelationshipStoreV2(typeRegistry, deleteDelegate, entityNotifier);
+        entityStore       = new AtlasEntityStoreV2(atlasGraph, deleteDelegate, typeRegistry, mockChangeNotifier, graphMapper);
+        relationshipStore = new AtlasRelationshipStoreV2(atlasGraph, typeRegistry, deleteDelegate, entityNotifier);
 
         RequestContext.clear();
         RequestContext.get().setUser(TestUtilsV2.TEST_USER, null);

repository/src/test/java/org/apache/atlas/repository/store/graph/v2/ClassificationAssociatorTest.java

@@ -20,15 +20,13 @@ package org.apache.atlas.repository.store.graph.v2;
 
 import com.fasterxml.jackson.core.type.TypeReference;
 import org.apache.atlas.exception.AtlasBaseException;
-import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasEntityHeaders;
 import org.apache.atlas.model.typedef.AtlasEntityDef;
-import org.apache.atlas.model.typedef.AtlasStructDef;
 import org.apache.atlas.repository.audit.EntityAuditRepository;
+import org.apache.atlas.repository.graphdb.AtlasGraph;
 import org.apache.atlas.repository.store.graph.AtlasEntityStore;
 import org.apache.atlas.type.AtlasEntityType;
-import org.apache.atlas.type.AtlasStructType;
 import org.apache.atlas.type.AtlasTypeRegistry;
 import org.apache.atlas.utils.AtlasJson;
 import org.apache.atlas.utils.TestResourceFileUtils;
@@ -37,7 +35,6 @@ import org.elasticsearch.common.util.CollectionUtils;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
-import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -65,13 +62,13 @@ public class ClassificationAssociatorTest {
     private class ClassificationAssociatorUpdaterForSpy extends ClassificationAssociator.Updater {
         private final String entityFileName;
 
-        public ClassificationAssociatorUpdaterForSpy(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
-            super(typeRegistry, entitiesStore);
+        public ClassificationAssociatorUpdaterForSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
+            super(atlasGraph, typeRegistry, entitiesStore);
             this.entityFileName = StringUtils.EMPTY;
         }
 
-        public ClassificationAssociatorUpdaterForSpy(AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore, String entityFileName) {
-            super(typeRegistry, entitiesStore);
+        public ClassificationAssociatorUpdaterForSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore, String entityFileName) {
+            super(atlasGraph, typeRegistry, entitiesStore);
             this.entityFileName = entityFileName;
         }
 
@@ -142,7 +139,8 @@ public class ClassificationAssociatorTest {
         AtlasTypeRegistry typeRegistry = mock(AtlasTypeRegistry.class);
         when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(null);
 
-        ClassificationAssociator.Updater updater = new ClassificationAssociator.Updater(typeRegistry, entitiesStore);
+        AtlasGraph atlasGraph = mock(AtlasGraph.class);
+        ClassificationAssociator.Updater updater = new ClassificationAssociator.Updater(atlasGraph, typeRegistry, entitiesStore);
         String summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
 
         assertTrue(summary.contains("hive_"));
@@ -155,11 +153,12 @@ public class ClassificationAssociatorTest {
         AtlasEntityType hiveTable = mock(AtlasEntityType.class);
         AtlasEntityStore entitiesStore = mock(AtlasEntityStore.class);
         AtlasTypeRegistry typeRegistry = mock(AtlasTypeRegistry.class);
+        AtlasGraph atlasGraph = mock(AtlasGraph.class);
 
         when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(hiveTable);
         when(hiveTable.getTypeName()).thenReturn("hive_column");
 
-        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(typeRegistry, entitiesStore);
+        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(atlasGraph, typeRegistry, entitiesStore);
         String summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
 
         TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
@@ -178,8 +177,9 @@ public class ClassificationAssociatorTest {
         AtlasTypeRegistry typeRegistry = new AtlasTypeRegistry();
         AtlasTypeRegistry.AtlasTransientTypeRegistry ttr = typeRegistry.lockTypeRegistryForUpdate();
         ttr.addTypes(CollectionUtils.newSingletonArrayList(ed));
+        AtlasGraph atlasGraph = mock(AtlasGraph.class);
 
-        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(ttr, entitiesStore, "col-entity-PII");
+        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(atlasGraph, ttr, entitiesStore, "col-entity-PII");
         String summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
 
         TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
@@ -233,11 +233,12 @@ public class ClassificationAssociatorTest {
         AtlasEntityType hiveTable = mock(AtlasEntityType.class);
         AtlasEntityStore entitiesStore = mock(AtlasEntityStore.class);
         AtlasTypeRegistry typeRegistry = mock(AtlasTypeRegistry.class);
+        AtlasGraph atlasGraph = mock(AtlasGraph.class);
 
         when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(hiveTable);
         when(hiveTable.getTypeName()).thenReturn("hive_column");
 
-        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(typeRegistry, entitiesStore, entityFileName);
+        ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(atlasGraph, typeRegistry, entitiesStore, entityFileName);
         String summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
 
         TypeReference<String[]> typeReference = new TypeReference<String[]>() {};