From 15d7bf7ce73844a17e55e4a7c73be3deb0d8cbd6 Mon Sep 17 00:00:00 2001
From: Andrew Shvayka <ashvayka@thingsboard.io>
Date: Thu, 28 Jun 2018 18:54:49 +0300
Subject: [PATCH] Added RelationTypeGroup to relation Query API

---
 .../data/relation/RelationsSearchParameters.java  |  6 ++++++
 .../server/dao/relation/BaseRelationService.java  | 15 ++++++++-------
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/relation/RelationsSearchParameters.java b/common/data/src/main/java/org/thingsboard/server/common/data/relation/RelationsSearchParameters.java
index 2e976f20a6..beccc04a48 100644
--- a/common/data/src/main/java/org/thingsboard/server/common/data/relation/RelationsSearchParameters.java
+++ b/common/data/src/main/java/org/thingsboard/server/common/data/relation/RelationsSearchParameters.java
@@ -33,13 +33,19 @@ public class RelationsSearchParameters {
     private UUID rootId;
     private EntityType rootType;
     private EntitySearchDirection direction;
+    private RelationTypeGroup relationTypeGroup;
     private int maxLevel = 1;
 
     public RelationsSearchParameters(EntityId entityId, EntitySearchDirection direction, int maxLevel) {
+        this(entityId, direction, maxLevel, RelationTypeGroup.COMMON);
+    }
+
+    public RelationsSearchParameters(EntityId entityId, EntitySearchDirection direction, int maxLevel, RelationTypeGroup relationTypeGroup) {
         this.rootId = entityId.getId();
         this.rootType = entityId.getEntityType();
         this.direction = direction;
         this.maxLevel = maxLevel;
+        this.relationTypeGroup = relationTypeGroup;
     }
 
     public EntityId getEntityId() {
diff --git a/dao/src/main/java/org/thingsboard/server/dao/relation/BaseRelationService.java b/dao/src/main/java/org/thingsboard/server/dao/relation/BaseRelationService.java
index ea29f6fa77..59668c893b 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/relation/BaseRelationService.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/relation/BaseRelationService.java
@@ -402,7 +402,7 @@ public class BaseRelationService implements RelationService {
         int maxLvl = params.getMaxLevel() > 0 ? params.getMaxLevel() : Integer.MAX_VALUE;
 
         try {
-            ListenableFuture<Set<EntityRelation>> relationSet = findRelationsRecursively(params.getEntityId(), params.getDirection(), maxLvl, new ConcurrentHashMap<>());
+            ListenableFuture<Set<EntityRelation>> relationSet = findRelationsRecursively(params.getEntityId(), params.getDirection(), params.getRelationTypeGroup(), maxLvl, new ConcurrentHashMap<>());
             return Futures.transform(relationSet, input -> {
                 List<EntityRelation> relations = new ArrayList<>();
                 if (filters == null || filters.isEmpty()) {
@@ -518,14 +518,15 @@ public class BaseRelationService implements RelationService {
         }
     }
 
-    private ListenableFuture<Set<EntityRelation>> findRelationsRecursively(final EntityId rootId, final EntitySearchDirection direction, int lvl,
+    private ListenableFuture<Set<EntityRelation>> findRelationsRecursively(final EntityId rootId, final EntitySearchDirection direction,
+                                                                           RelationTypeGroup relationTypeGroup, int lvl,
                                                                            final ConcurrentHashMap<EntityId, Boolean> uniqueMap) throws Exception {
         if (lvl == 0) {
             return Futures.immediateFuture(Collections.emptySet());
         }
         lvl--;
         //TODO: try to remove this blocking operation
-        Set<EntityRelation> children = new HashSet<>(findRelations(rootId, direction).get());
+        Set<EntityRelation> children = new HashSet<>(findRelations(rootId, direction, relationTypeGroup).get());
         Set<EntityId> childrenIds = new HashSet<>();
         for (EntityRelation childRelation : children) {
             log.trace("Found Relation: {}", childRelation);
@@ -544,7 +545,7 @@ public class BaseRelationService implements RelationService {
         }
         List<ListenableFuture<Set<EntityRelation>>> futures = new ArrayList<>();
         for (EntityId entityId : childrenIds) {
-            futures.add(findRelationsRecursively(entityId, direction, lvl, uniqueMap));
+            futures.add(findRelationsRecursively(entityId, direction, relationTypeGroup, lvl, uniqueMap));
         }
         //TODO: try to remove this blocking operation
         List<Set<EntityRelation>> relations = Futures.successfulAsList(futures).get();
@@ -552,12 +553,12 @@ public class BaseRelationService implements RelationService {
         return Futures.immediateFuture(children);
     }
 
-    private ListenableFuture<List<EntityRelation>> findRelations(final EntityId rootId, final EntitySearchDirection direction) {
+    private ListenableFuture<List<EntityRelation>> findRelations(final EntityId rootId, final EntitySearchDirection direction, RelationTypeGroup relationTypeGroup) {
         ListenableFuture<List<EntityRelation>> relations;
         if (direction == EntitySearchDirection.FROM) {
-            relations = findByFromAsync(rootId, RelationTypeGroup.COMMON);
+            relations = findByFromAsync(rootId, relationTypeGroup);
         } else {
-            relations = findByToAsync(rootId, RelationTypeGroup.COMMON);
+            relations = findByToAsync(rootId, relationTypeGroup);
         }
         return relations;
     }