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);

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);
         }