Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

PLT-573 | Beta | Add removeSupport in bulk API #2745

Merged
merged 12 commits into from
Jan 11, 2024
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,12 @@ public Object getAppendRelationshipAttribute(String name) {
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;

Expand Down Expand Up @@ -404,7 +410,24 @@ public void setRemoveRelationshipAttributes(Map<String, Object> removeRelationsh
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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ public final class EntityGraphDiscoveryContext {
private final Set<String> localGuids = new HashSet<>();

private boolean isAppendRelationshipAttributeVisited;
private boolean isRemoveRelationshipAttributeVisited;

public EntityGraphDiscoveryContext(AtlasTypeRegistry typeRegistry, EntityStream entityStream) {
this.typeRegistry = typeRegistry;
Expand Down Expand Up @@ -164,4 +165,12 @@ public boolean 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
Expand Up @@ -376,15 +376,28 @@ 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);
} 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());
visitAttribute(attribute.getAttributeType(), attrVal, entity.getGuid());

visitedAttributes.add(attrName);
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);
}
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1708,6 +1708,10 @@ private EntityMutationContext preCreateOrUpdate(EntityStream entityStream, Entit
context.setUpdatedWithRelationshipAttributes(entity);
}

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

if (isRestoreRequested) {
Status currStatus = AtlasGraphUtilsV2.getState(vertex);
if (currStatus == Status.DELETED) {
Expand Down
Loading
Loading