Skip to content

A
atlas
Commit 2b38550b by Ashutosh Mestry
Options

ATLAS-3737: Refactoring: Support multiple instances of AtlasGraph

parent 40fa099c
Showing with 538 additions and 347 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/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[]>() {};
......
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