Skip to content

A
atlas
Commit 2b38550b by Ashutosh Mestry
Options

ATLAS-3737: Refactoring: Support multiple instances of AtlasGraph

parent 40fa099c
Showing with 392 additions and 234 deletions
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/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/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[]>() {};
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment