ATLAS-2619: Add query param in relationship GET API to get extended info about referred entities

intg/src/main/java/org/apache/atlas/model/instance/AtlasRelationship.java

@@ -32,6 +32,7 @@ import javax.xml.bind.annotation.XmlRootElement;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Date;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -407,4 +408,89 @@ public class AtlasRelationship extends AtlasStruct implements Serializable {
     public String toString() {
         return toString(new StringBuilder()).toString();
     }
+
+    @JsonAutoDetect(getterVisibility=PUBLIC_ONLY, setterVisibility=PUBLIC_ONLY, fieldVisibility=NONE)
+    @JsonSerialize(include=JsonSerialize.Inclusion.NON_NULL)
+    @JsonIgnoreProperties(ignoreUnknown=true)
+    @XmlRootElement
+    @XmlAccessorType(XmlAccessType.PROPERTY)
+    public static class AtlasRelationshipWithExtInfo implements Serializable {
+        private AtlasRelationship              relationship;
+        private Map<String, AtlasEntityHeader> referredEntities;
+
+        public AtlasRelationshipWithExtInfo() {
+        }
+
+        public AtlasRelationshipWithExtInfo(AtlasRelationship relationship) {
+            setRelationship(relationship);
+        }
+
+        public AtlasRelationship getRelationship() {
+            return relationship;
+        }
+
+        public void setRelationship(AtlasRelationship relationship) {
+            this.relationship = relationship;
+        }
+
+        public Map<String, AtlasEntityHeader> getReferredEntities() {
+            return referredEntities;
+        }
+
+        public void setReferredEntities(Map<String, AtlasEntityHeader> referredEntities) {
+            this.referredEntities = referredEntities;
+        }
+
+        public boolean referredEntitiesContains(String guid) {
+            return (referredEntities != null) ? referredEntities.containsKey(guid) : false;
+        }
+
+        @JsonIgnore
+        public final void addReferredEntity(String guid, AtlasEntityHeader entityHeader) {
+            Map<String, AtlasEntityHeader> r = this.referredEntities;
+
+            if (r == null) {
+                r = new HashMap<>();
+
+                this.referredEntities = r;
+            }
+
+            if (guid != null) {
+                r.put(guid, entityHeader);
+            }
+        }
+
+        @JsonIgnore
+        public final AtlasEntityHeader removeReferredEntity(String guid) {
+            Map<String, AtlasEntityHeader> r = this.referredEntities;
+
+            return r != null && guid != null ? r.remove(guid) : null;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) { return true; }
+            if (o == null || getClass() != o.getClass()) { return false; }
+            if (!super.equals(o)) { return false; }
+            AtlasRelationshipWithExtInfo that = (AtlasRelationshipWithExtInfo) o;
+
+            return Objects.equals(relationship, that.relationship) &&
+                   Objects.equals(referredEntities, that.referredEntities);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(super.hashCode(), relationship, referredEntities);
+        }
+
+        @Override
+        public String toString() {
+            final StringBuilder sb = new StringBuilder("AtlasRelationshipWithExtInfo{");
+            sb.append("relationship=").append(relationship);
+            sb.append(", referredEntities=").append(referredEntities);
+            sb.append('}');
+
+            return sb.toString();
+        }
+    }
 }
\ No newline at end of file

repository/src/main/java/org/apache/atlas/repository/store/graph/AtlasRelationshipStore.java

@@ -19,6 +19,7 @@ package org.apache.atlas.repository.store.graph;
 
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasRelationship;
+import org.apache.atlas.model.instance.AtlasRelationship.AtlasRelationshipWithExtInfo;
 import org.apache.atlas.repository.graphdb.AtlasEdge;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 
@@ -48,6 +49,13 @@ public interface AtlasRelationshipStore {
      */
     AtlasRelationship getById(String guid) throws AtlasBaseException;
 
+    /**
+     * Retrieve a relationship instance and its referred entities using guid.
+     * @param guid relationship instance guid
+     * @return AtlasRelationship
+     */
+    AtlasRelationshipWithExtInfo getExtInfoById(String guid) throws AtlasBaseException;
+
 
     AtlasEdge getOrCreate(AtlasVertex end1Vertex, AtlasVertex end2Vertex, AtlasRelationship relationship) throws AtlasBaseException;
 

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

@@ -19,7 +19,6 @@
 package org.apache.atlas.repository.store.graph.v1;
 
 import org.apache.atlas.AtlasErrorCode;
-import org.apache.atlas.RequestContextV1;
 import org.apache.atlas.annotation.GraphTransaction;
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.TypeCategory;
@@ -27,6 +26,7 @@ import org.apache.atlas.model.instance.AtlasClassification;
 import org.apache.atlas.model.instance.AtlasEntity.Status;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasRelationship;
+import org.apache.atlas.model.instance.AtlasRelationship.AtlasRelationshipWithExtInfo;
 import org.apache.atlas.model.typedef.AtlasRelationshipDef;
 import org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags;
 import org.apache.atlas.model.typedef.AtlasRelationshipEndDef;
@@ -204,11 +204,8 @@ public class AtlasRelationshipStoreV1 implements AtlasRelationshipStore {
             LOG.debug("==> getById({})", guid);
         }
 
-        AtlasRelationship ret;
-
         AtlasEdge         edge = graphHelper.getEdgeForGUID(guid);
-
-        ret = entityRetriever.mapEdgeToAtlasRelationship(edge);
+        AtlasRelationship ret  = entityRetriever.mapEdgeToAtlasRelationship(edge);
 
         if (LOG.isDebugEnabled()) {
             LOG.debug("<== getById({}): {}", guid, ret);
@@ -219,6 +216,23 @@ public class AtlasRelationshipStoreV1 implements AtlasRelationshipStore {
 
     @Override
     @GraphTransaction
+    public AtlasRelationshipWithExtInfo getExtInfoById(String guid) throws AtlasBaseException {
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("==> getExtInfoById({})", guid);
+        }
+
+        AtlasEdge                    edge = graphHelper.getEdgeForGUID(guid);
+        AtlasRelationshipWithExtInfo ret  = entityRetriever.mapEdgeToAtlasRelationshipWithExtInfo(edge);
+
+        if (LOG.isDebugEnabled()) {
+            LOG.debug("<== getExtInfoById({}): {}", guid, ret);
+        }
+
+        return ret;
+    }
+
+    @Override
+    @GraphTransaction
     public void deleteById(String guid) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("==> deleteById({})", guid);

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

@@ -33,9 +33,9 @@ import org.apache.atlas.model.instance.AtlasEntityHeader;
 import org.apache.atlas.model.instance.AtlasObjectId;
 import org.apache.atlas.model.instance.AtlasRelatedObjectId;
 import org.apache.atlas.model.instance.AtlasRelationship;
+import org.apache.atlas.model.instance.AtlasRelationship.AtlasRelationshipWithExtInfo;
 import org.apache.atlas.model.instance.AtlasStruct;
 import org.apache.atlas.model.typedef.AtlasRelationshipDef;
-import org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags;
 import org.apache.atlas.model.typedef.AtlasRelationshipEndDef;
 import org.apache.atlas.model.typedef.AtlasStructDef.AtlasAttributeDef;
 import org.apache.atlas.repository.Constants;
@@ -45,7 +45,6 @@ import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
 import org.apache.atlas.repository.graphdb.AtlasElement;
 import org.apache.atlas.repository.graphdb.AtlasVertex;
 import org.apache.atlas.type.AtlasArrayType;
-import org.apache.atlas.type.AtlasClassificationType;
 import org.apache.atlas.type.AtlasEntityType;
 import org.apache.atlas.type.AtlasMapType;
 import org.apache.atlas.type.AtlasRelationshipType;
@@ -72,36 +71,45 @@ import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
-import static org.apache.atlas.glossary.GlossaryUtils.*;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_CONFIDENCE;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_CREATED_BY;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_DESCRIPTION;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_EXPRESSION;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_SOURCE;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_STATUS;
+import static org.apache.atlas.glossary.GlossaryUtils.TERM_ASSIGNMENT_ATTR_STEWARD;
 import static org.apache.atlas.model.instance.AtlasClassification.PropagationState.ACTIVE;
 import static org.apache.atlas.model.instance.AtlasClassification.PropagationState.DELETED;
-import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.*;
-import static org.apache.atlas.model.typedef.AtlasRelationshipDef.PropagateTags.ONE_TO_TWO;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BIGDECIMAL;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BIGINTEGER;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BOOLEAN;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_BYTE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_DATE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_DOUBLE;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_FLOAT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_INT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_LONG;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_SHORT;
+import static org.apache.atlas.model.typedef.AtlasBaseTypeDef.ATLAS_TYPE_STRING;
 import static org.apache.atlas.repository.Constants.CLASSIFICATION_ENTITY_GUID;
 import static org.apache.atlas.repository.Constants.CLASSIFICATION_LABEL;
 import static org.apache.atlas.repository.Constants.CLASSIFICATION_VALIDITY_PERIODS_KEY;
 import static org.apache.atlas.repository.Constants.TERM_ASSIGNMENT_LABEL;
 import static org.apache.atlas.repository.graph.GraphHelper.EDGE_LABEL_PREFIX;
-import static org.apache.atlas.repository.graph.GraphHelper.addToPropagatedTraitNames;
 import static org.apache.atlas.repository.graph.GraphHelper.getAdjacentEdgesByLabel;
 import static org.apache.atlas.repository.graph.GraphHelper.getAllClassificationEdges;
 import static org.apache.atlas.repository.graph.GraphHelper.getAllTraitNames;
-import static org.apache.atlas.repository.graph.GraphHelper.getAssociatedEntityVertex;
 import static org.apache.atlas.repository.graph.GraphHelper.getBlockedClassificationIds;
-import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdge;
 import static org.apache.atlas.repository.graph.GraphHelper.getClassificationEdgeState;
 import static org.apache.atlas.repository.graph.GraphHelper.getClassificationVertices;
 import static org.apache.atlas.repository.graph.GraphHelper.getGuid;
 import static org.apache.atlas.repository.graph.GraphHelper.getIncomingEdgesByLabel;
 import static org.apache.atlas.repository.graph.GraphHelper.getOutGoingEdgesByLabel;
 import static org.apache.atlas.repository.graph.GraphHelper.getPropagateTags;
-import static org.apache.atlas.repository.graph.GraphHelper.getPropagatedClassificationEdge;
-import static org.apache.atlas.repository.graph.GraphHelper.getPropagationEnabledClassificationVertices;
 import static org.apache.atlas.repository.graph.GraphHelper.getRelationshipGuid;
 import static org.apache.atlas.repository.graph.GraphHelper.getTypeName;
 import static org.apache.atlas.repository.graph.GraphHelper.isPropagatedClassificationEdge;
 import static org.apache.atlas.repository.graph.GraphHelper.isPropagationEnabled;
-import static org.apache.atlas.repository.graph.GraphHelper.removeFromPropagatedTraitNames;
 import static org.apache.atlas.repository.store.graph.v1.AtlasGraphUtilsV1.getIdFromVertex;
 import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection;
 import static org.apache.atlas.type.AtlasStructType.AtlasAttribute.AtlasRelationshipEdgeDirection.BOTH;
@@ -966,20 +974,36 @@ public final class EntityGraphRetriever {
     }
 
     public AtlasRelationship mapEdgeToAtlasRelationship(AtlasEdge edge) throws AtlasBaseException {
-        AtlasRelationship ret = new AtlasRelationship();
+        return mapEdgeToAtlasRelationship(edge, false).getRelationship();
+    }
+
+    public AtlasRelationshipWithExtInfo mapEdgeToAtlasRelationshipWithExtInfo(AtlasEdge edge) throws AtlasBaseException {
+        return mapEdgeToAtlasRelationship(edge, true);
+    }
 
-        mapSystemAttributes(edge, ret);
+    public AtlasRelationshipWithExtInfo mapEdgeToAtlasRelationship(AtlasEdge edge, boolean extendedInfo) throws AtlasBaseException {
+        AtlasRelationshipWithExtInfo ret = new AtlasRelationshipWithExtInfo();
+
+        mapSystemAttributes(edge, ret, extendedInfo);
 
         mapAttributes(edge, ret);
 
         return ret;
     }
 
-    private AtlasRelationship mapSystemAttributes(AtlasEdge edge, AtlasRelationship relationship) throws AtlasBaseException {
+    private AtlasRelationshipWithExtInfo mapSystemAttributes(AtlasEdge edge, AtlasRelationshipWithExtInfo relationshipWithExtInfo, boolean extendedInfo) throws AtlasBaseException {
         if (LOG.isDebugEnabled()) {
             LOG.debug("Mapping system attributes for relationship");
         }
 
+        AtlasRelationship relationship = relationshipWithExtInfo.getRelationship();
+
+        if (relationship == null) {
+            relationship = new AtlasRelationship();
+
+            relationshipWithExtInfo.setRelationship(relationship);
+        }
+
         relationship.setGuid(getRelationshipGuid(edge));
         relationship.setTypeName(getTypeName(edge));
 
@@ -1007,25 +1031,38 @@ public final class EntityGraphRetriever {
         relationship.setLabel(edge.getLabel());
         relationship.setPropagateTags(getPropagateTags(edge));
 
+        if (extendedInfo) {
+            addToReferredEntities(relationshipWithExtInfo, end1Vertex);
+            addToReferredEntities(relationshipWithExtInfo, end2Vertex);
+        }
+
         // set propagated and blocked propagated classifications
-        readClassificationsFromEdge(edge, relationship);
+        readClassificationsFromEdge(edge, relationshipWithExtInfo, extendedInfo);
 
-        return relationship;
+        return relationshipWithExtInfo;
     }
 
-    private void readClassificationsFromEdge(AtlasEdge edge, AtlasRelationship relationship) throws AtlasBaseException {
+    private void readClassificationsFromEdge(AtlasEdge edge, AtlasRelationshipWithExtInfo relationshipWithExtInfo, boolean extendedInfo) throws AtlasBaseException {
         List<AtlasVertex>         classificationVertices    = getClassificationVertices(edge);
         List<String>              blockedClassificationIds  = getBlockedClassificationIds(edge);
         List<AtlasClassification> propagatedClassifications = new ArrayList<>();
         List<AtlasClassification> blockedClassifications    = new ArrayList<>();
+        AtlasRelationship         relationship              = relationshipWithExtInfo.getRelationship();
 
         for (AtlasVertex classificationVertex : classificationVertices) {
             String              classificationId = classificationVertex.getIdForDisplay();
+            AtlasClassification classification   = toAtlasClassification(classificationVertex);
+            String              entityGuid       = classification.getEntityGuid();
 
             if (blockedClassificationIds.contains(classificationId)) {
-                blockedClassifications.add(toAtlasClassification(classificationVertex));
+                blockedClassifications.add(classification);
             } else {
-                propagatedClassifications.add(toAtlasClassification(classificationVertex));
+                propagatedClassifications.add(classification);
+            }
+
+            // add entity headers to referred entities
+            if (extendedInfo) {
+                addToReferredEntities(relationshipWithExtInfo, entityGuid);
             }
         }
 
@@ -1033,7 +1070,22 @@ public final class EntityGraphRetriever {
         relationship.setBlockedPropagatedClassifications(blockedClassifications);
     }
 
-    private void mapAttributes(AtlasEdge edge, AtlasRelationship relationship) throws AtlasBaseException {
+    private void addToReferredEntities(AtlasRelationshipWithExtInfo relationshipWithExtInfo, String guid) throws AtlasBaseException {
+        if (!relationshipWithExtInfo.referredEntitiesContains(guid)) {
+            addToReferredEntities(relationshipWithExtInfo, getEntityVertex(guid));
+        }
+    }
+
+    private void addToReferredEntities(AtlasRelationshipWithExtInfo relationshipWithExtInfo, AtlasVertex entityVertex) throws AtlasBaseException {
+        String entityGuid = getGuid(entityVertex);
+
+        if (!relationshipWithExtInfo.referredEntitiesContains(entityGuid)) {
+            relationshipWithExtInfo.addReferredEntity(entityGuid, toAtlasEntityHeader(entityVertex));
+        }
+    }
+
+    private void mapAttributes(AtlasEdge edge, AtlasRelationshipWithExtInfo relationshipWithExtInfo) throws AtlasBaseException {
+        AtlasRelationship relationship = relationshipWithExtInfo.getRelationship();
         AtlasType         objType      = typeRegistry.getType(relationship.getTypeName());
 
         if (!(objType instanceof AtlasRelationshipType)) {

webapp/src/main/java/org/apache/atlas/web/rest/RelationshipREST.java

@@ -20,6 +20,7 @@ package org.apache.atlas.web.rest;
 
 import org.apache.atlas.exception.AtlasBaseException;
 import org.apache.atlas.model.instance.AtlasRelationship;
+import org.apache.atlas.model.instance.AtlasRelationship.AtlasRelationshipWithExtInfo;
 import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
 import org.apache.atlas.utils.AtlasPerfTracer;
 import org.apache.atlas.web.util.Servlets;
@@ -30,12 +31,14 @@ import javax.inject.Inject;
 import javax.inject.Singleton;
 import javax.ws.rs.Consumes;
 import javax.ws.rs.DELETE;
+import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
 import javax.ws.rs.POST;
 import javax.ws.rs.PUT;
 import javax.ws.rs.Path;
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
+import javax.ws.rs.QueryParam;
 
 /**
  * REST interface for entity relationships.
@@ -102,18 +105,27 @@ public class RelationshipREST {
     @Path("/guid/{guid}")
     @Consumes(Servlets.JSON_MEDIA_TYPE)
     @Produces(Servlets.JSON_MEDIA_TYPE)
-    public AtlasRelationship getById(@PathParam("guid") String guid) throws AtlasBaseException {
+    public AtlasRelationshipWithExtInfo getById(@PathParam("guid") String guid,
+                                                @QueryParam("extendedInfo") @DefaultValue("false") boolean extendedInfo)
+                                                throws AtlasBaseException {
         Servlets.validateQueryParamLength("guid", guid);
 
         AtlasPerfTracer perf = null;
 
+        AtlasRelationshipWithExtInfo ret;
+
         try {
             if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
                 perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "RelationshipREST.getById(" + guid + ")");
             }
 
-            return relationshipStore.getById(guid);
+            if (extendedInfo) {
+                ret = relationshipStore.getExtInfoById(guid);
+            } else {
+                ret = new AtlasRelationshipWithExtInfo(relationshipStore.getById(guid));
+            }
 
+            return ret;
         } finally {
             AtlasPerfTracer.log(perf);
         }