Skip to content

Commit

Permalink
Merge pull request #2744 from atlanhq/feat/removeRelationshipAttribute
Browse files Browse the repository at this point in the history
PLT-573 | Add removeSupport in bulk API
aarshi0301 authored Jan 17, 2024

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
2 parents a055d01 + 71dbaeb commit 7194492
Showing 8 changed files with 505 additions and 21 deletions.
Original file line number Diff line number Diff line change
@@ -92,6 +92,8 @@ public enum Status { ACTIVE, DELETED, PURGED }
private Boolean starred = null;

private Map<String, Object> relationshipAttributes;
private Map<String, Object> appendRelationshipAttributes;
private Map<String, Object> removeRelationshipAttributes;
private List<AtlasClassification> classifications;
private List<AtlasTermAssignmentHeader> meanings;
private Map<String, String> customAttributes;
@@ -230,6 +232,8 @@ public AtlasEntity(AtlasEntity other) {
setBusinessAttributes(other.getBusinessAttributes());
setLabels(other.getLabels());
setPendingTasks(other.getPendingTasks());
setAppendRelationshipAttributes(other.getAppendRelationshipAttributes());
setRemoveRelationshipAttributes(other.getRemoveRelationshipAttributes());
}
}

@@ -354,12 +358,76 @@ public Object getRelationshipAttribute(String name) {
return a != null ? a.get(name) : null;
}

public Object getAppendRelationshipAttribute(String name) {
Map<String, Object> a = this.appendRelationshipAttributes;

return a != null ? a.get(name) : null;
}

public Object getRemoveRelationshipAttribute(String name) {
Map<String, Object> a = this.removeRelationshipAttributes;

return a != null ? a.get(name) : null;
}

public boolean hasRelationshipAttribute(String name) {
Map<String, Object> r = this.relationshipAttributes;

return r != null ? r.containsKey(name) : false;
}

public boolean hasAppendRelationshipAttribute(String name) {
Map<String, Object> r = this.appendRelationshipAttributes;

return r != null ? r.containsKey(name) : false;
}

public Map<String, Object> getAppendRelationshipAttributes() {
return appendRelationshipAttributes;
}

public void setAppendRelationshipAttributes(Map<String, Object> appendRelationshipAttributes) {
this.appendRelationshipAttributes = appendRelationshipAttributes;
}

public void setAppendRelationshipAttribute(String name, Object value) {
Map<String, Object> r = this.appendRelationshipAttributes;

if (r != null) {
r.put(name, value);
} else {
r = new HashMap<>();
r.put(name, value);

this.appendRelationshipAttributes = r;
}
}
public Map<String, Object> getRemoveRelationshipAttributes() {
return removeRelationshipAttributes;
}

public void setRemoveRelationshipAttributes(Map<String, Object> removeRelationshipAttributes) {
this.removeRelationshipAttributes = removeRelationshipAttributes;
}

public boolean hasRemoveRelationshipAttribute(String name) {
Map<String, Object> r = this.removeRelationshipAttributes;

return r != null ? r.containsKey(name) : false;
}

public void setRemoveRelationshipAttribute(String name, Object value) {
Map<String, Object> r = this.removeRelationshipAttributes;

if (r != null) {
r.put(name, value);
} else {
r = new HashMap<>();
r.put(name, value);

this.removeRelationshipAttributes = r;
}
}
public Map<String, String> getCustomAttributes() {
return customAttributes;
}
Original file line number Diff line number Diff line change
@@ -70,6 +70,7 @@ public interface AtlasRelationshipStore {

AtlasEdge getOrCreate(AtlasVertex end1Vertex, AtlasVertex end2Vertex, AtlasRelationship relationship) throws AtlasBaseException;

AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException;

/**
* Retrieve a relationship if it exists or creates a new relationship instance.
Original file line number Diff line number Diff line change
@@ -44,6 +44,9 @@ public final class EntityGraphDiscoveryContext {
private final Map<AtlasObjectId, AtlasVertex> resolvedIdsByUniqAttribs = new HashMap<>();
private final Set<String> localGuids = new HashSet<>();

private boolean isAppendRelationshipAttributeVisited;
private boolean isRemoveRelationshipAttributeVisited;

public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
this.typeRegistry = typeRegistry;
this.entityStream = entityStream;
@@ -154,4 +157,20 @@ public void cleanUp() {
resolvedIdsByUniqAttribs.clear();
localGuids.clear();
}

public boolean isAppendRelationshipAttributeVisited() {
return isAppendRelationshipAttributeVisited;
}

public void setAppendRelationshipAttributeVisited(boolean appendRelationshipAttributeVisited) {
isAppendRelationshipAttributeVisited = appendRelationshipAttributeVisited;
}

public boolean isRemoveRelationshipAttributeVisited() {
return isRemoveRelationshipAttributeVisited;
}

public void setRemoveRelationshipAttributeVisited(boolean removeRelationshipAttributeVisited) {
isRemoveRelationshipAttributeVisited = removeRelationshipAttributeVisited;
}
}
Original file line number Diff line number Diff line change
@@ -358,6 +358,8 @@ private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity,
for (String attrName : entityType.getRelationshipAttributes().keySet()) {

// if attribute is not in 'relationshipAttributes', try 'attributes'
// appendRelationshipAttribute will be ignored if same attribute is present
// in relationshipAttribute
if (entity.hasRelationshipAttribute(attrName)) {
Object attrVal = entity.getRelationshipAttribute(attrName);
String relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
@@ -374,10 +376,31 @@ private void visitRelationships(AtlasEntityType entityType, AtlasEntity entity,
visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());

visitedAttributes.add(attrName);
} else {
if (entity.hasAppendRelationshipAttribute(attrName)) {
discoveryContext.setAppendRelationshipAttributeVisited(true);
Object attrVal = entity.getAppendRelationshipAttribute(attrName);
String relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
AtlasAttribute attribute = entityType.getRelationshipAttribute(attrName, relationshipType);

visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());

visitedAttributes.add(attrName);
}

if (entity.hasRemoveRelationshipAttribute(attrName)) {
discoveryContext.setRemoveRelationshipAttributeVisited(true);
Object attrVal = entity.getRemoveRelationshipAttribute(attrName);
String relationshipType = AtlasEntityUtil.getRelationshipType(attrVal);
AtlasAttribute attribute = entityType.getRelationshipAttribute(attrName, relationshipType);

visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());

visitedAttributes.add(attrName);
}
}
}
}

void visitStruct(AtlasStructType structType, AtlasStruct struct, String referringEntityGuid) throws AtlasBaseException {
for (AtlasAttribute attribute : structType.getAllAttributes().values()) {
AtlasType attrType = attribute.getAttributeType();
Original file line number Diff line number Diff line change
@@ -1438,7 +1438,6 @@ private EntityMutationResponse createOrUpdate(EntityStream entityStream, boolean
}
}
}

// for existing entities, skip update if incoming entity doesn't have any change
if (CollectionUtils.isNotEmpty(context.getUpdatedEntities())) {
MetricRecorder checkForUnchangedEntities = RequestContext.get().startMetricRecord("checkForUnchangedEntities");
@@ -1644,6 +1643,14 @@ private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, Entit
String entityActiveKey = Status.ACTIVE.toString();
boolean isRestoreRequested = ((StringUtils.isNotEmpty(entityStateValue) && entityStateValue.equals(entityActiveKey)) || (StringUtils.isNotEmpty(entityStatusValue) && entityStatusValue.equals(entityActiveKey)));

if (discoveryContext.isAppendRelationshipAttributeVisited() && MapUtils.isNotEmpty(entity.getAppendRelationshipAttributes())) {
context.setUpdatedWithRelationshipAttributes(entity);
}

if (discoveryContext.isRemoveRelationshipAttributeVisited() && MapUtils.isNotEmpty(entity.getRemoveRelationshipAttributes())) {
context.setUpdatedWithRemoveRelationshipAttributes(entity);
}

if (isRestoreRequested) {
Status currStatus = AtlasGraphUtilsV2.getState(vertex);
if (currStatus == Status.DELETED) {
Original file line number Diff line number Diff line change
@@ -397,7 +397,8 @@ public AtlasEdge getOrCreate(AtlasVertex end1Vertex, AtlasVertex end2Vertex, Atl
return ret;
}

private AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException {
@Override
public AtlasEdge getRelationship(AtlasVertex fromVertex, AtlasVertex toVertex, AtlasRelationship relationship) throws AtlasBaseException {
String relationshipLabel = getRelationshipEdgeLabel(fromVertex, toVertex, relationship.getTypeName());

return getRelationshipEdge(fromVertex, toVertex, relationshipLabel);

Large diffs are not rendered by default.

Original file line number Diff line number Diff line change
@@ -30,6 +30,8 @@ public class EntityMutationContext {
private final EntityGraphDiscoveryContext context;
private final List<AtlasEntity> entitiesCreated = new ArrayList<>();
private final List<AtlasEntity> entitiesUpdated = new ArrayList<>();
private final List<AtlasEntity> entitiesUpdatedWithAppendRelationshipAttribute = new ArrayList<>();
private final List<AtlasEntity> entitiesUpdatedWithRemoveRelationshipAttribute = new ArrayList<>();
private final Map<String, AtlasEntityType> entityVsType = new HashMap<>();
private final Map<String, AtlasVertex> entityVsVertex = new HashMap<>();
private final Map<String, String> guidAssignments = new HashMap<>();
@@ -57,6 +59,18 @@ public void addCreated(String internalGuid, AtlasEntity entity, AtlasEntityType
}
}

public void setUpdatedWithRelationshipAttributes(AtlasEntity entity){
entitiesUpdatedWithAppendRelationshipAttribute.add(entity);
}

public void setUpdatedWithRemoveRelationshipAttributes(AtlasEntity entity){
entitiesUpdatedWithRemoveRelationshipAttribute.add(entity);
}

public Collection<AtlasEntity> getEntitiesUpdatedWithRemoveRelationshipAttribute() {
return entitiesUpdatedWithRemoveRelationshipAttribute;
}

public void addUpdated(String internalGuid, AtlasEntity entity, AtlasEntityType type, AtlasVertex atlasVertex) {
if (!entityVsVertex.containsKey(internalGuid)) { // if the entity was already created/updated
entitiesUpdated.add(entity);
@@ -103,6 +117,10 @@ public Collection<AtlasEntity> getUpdatedEntities() {
return entitiesUpdated;
}

public Collection<AtlasEntity> getUpdatedEntitiesForAppendRelationshipAttribute() {
return entitiesUpdatedWithAppendRelationshipAttribute;
}

public Map<String, String> getGuidAssignments() {
return guidAssignments;
}

0 comments on commit 7194492

Please sign in to comment.