Fix #11944: API support to fetch count for mentions

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/CollectionDAO.java

@@ -116,8 +116,6 @@
 import org.openmetadata.schema.type.Include;
 import org.openmetadata.schema.type.Relationship;
 import org.openmetadata.schema.type.TagLabel;
-import org.openmetadata.schema.type.TaskStatus;
-import org.openmetadata.schema.type.ThreadType;
 import org.openmetadata.schema.type.UsageDetails;
 import org.openmetadata.schema.type.UsageStats;
 import org.openmetadata.schema.util.EntitiesCount;
@@ -1176,34 +1174,30 @@ int listCountThreadsByEntityLink(
     void update(@BindUUID("id") UUID id, @Bind("json") String json);
 
     @SqlQuery(
-        "SELECT entityLink, COUNT(id) count FROM field_relationship fr INNER JOIN thread_entity te ON fr.fromFQNHash=MD5(te.id) "
-            + "WHERE (:fqnPrefixHash IS NULL OR fr.toFQNHash LIKE CONCAT(:fqnPrefixHash, '.%') OR fr.toFQNHash=:fqnPrefixHash) AND "
-            + "(:toType IS NULL OR fr.toType like concat(:toType, '.%') OR fr.toType=:toType) AND fr.fromType = :fromType "
-            + "AND fr.relation = :relation AND te.resolved= :isResolved AND (:status IS NULL OR te.taskStatus = :status) "
-            + "AND (:type IS NULL OR te.type = :type) "
-            + "GROUP BY entityLink")
-    @RegisterRowMapper(CountFieldMapper.class)
+        "SELECT te.entityLink, te.type, te.taskStatus, COUNT(id) count FROM thread_entity te "
+            + " where entityId = :entityId  OR "
+            + " MD5(id) in (SELECT fromFQNHash FROM field_relationship WHERE "
+            + "(:fqnPrefixHash IS NULL OR toFQNHash LIKE CONCAT(:fqnPrefixHash, '.%') OR toFQNHash=:fqnPrefixHash) AND fromType='THREAD' AND "
+            + "(:toType IS NULL OR toType LIKE CONCAT(:toType, '.%') OR toType=:toType) AND relation= 3) "
+            + "GROUP BY te.type, te.taskStatus, entityLink")
+    @RegisterRowMapper(ThreadCountFieldMapper.class)
     List<List<String>> listCountByEntityLink(
+        @BindUUID("entityId") UUID entityId,
         @BindFQN("fqnPrefixHash") String fqnPrefixHash,
-        @Bind("fromType") String fromType,
-        @Bind("toType") String toType,
-        @Bind("relation") int relation,
-        @Bind("type") ThreadType type,
-        @Bind("status") TaskStatus status,
-        @Bind("isResolved") boolean isResolved);
+        @Bind("toType") String toType);
 
     @SqlQuery(
-        "SELECT entityLink, COUNT(id) count FROM thread_entity <condition> AND "
+        "SELECT te.type, te.taskStatus, COUNT(id) count FROM thread_entity te where "
             + "(entityId in (SELECT toId FROM entity_relationship WHERE "
             + "((fromEntity='user' AND fromId= :userId) OR "
             + "(fromEntity='team' AND fromId IN (<teamIds>))) AND relation=8) OR "
-            + "id in (SELECT toId FROM entity_relationship WHERE (fromEntity='user' AND fromId= :userId AND toEntity='THREAD' AND relation IN (1,2)))) "
-            + "GROUP BY entityLink")
-    @RegisterRowMapper(CountFieldMapper.class)
+            + "id in (SELECT toId FROM entity_relationship WHERE (fromEntity='user' AND fromId= :userId AND toEntity='THREAD' AND relation IN (1,2))) "
+            + " OR id in (SELECT toId FROM entity_relationship WHERE ((fromEntity='user' AND fromId= :userId) OR "
+            + "(fromEntity='team' AND fromId IN (<teamIds>))) AND relation=11)) "
+            + "GROUP BY te.type, te.taskStatus")
+    @RegisterRowMapper(OwnerCountFieldMapper.class)
     List<List<String>> listCountByOwner(
-        @BindUUID("userId") UUID userId,
-        @BindList("teamIds") List<String> teamIds,
-        @Define("condition") String condition);
+        @BindUUID("userId") UUID userId, @BindList("teamIds") List<String> teamIds);
 
     @SqlQuery(
         "SELECT json FROM thread_entity <condition> AND "
@@ -1292,10 +1286,22 @@ int listCountThreadsByMentions(
     @SqlQuery("select id from thread_entity where entityId = :entityId")
     List<String> findByEntityId(@Bind("entityId") String entityId);
 
-    class CountFieldMapper implements RowMapper<List<String>> {
+    class OwnerCountFieldMapper implements RowMapper<List<String>> {
+      @Override
+      public List<String> map(ResultSet rs, StatementContext ctx) throws SQLException {
+        return Arrays.asList(
+            rs.getString("type"), rs.getString("taskStatus"), rs.getString("count"));
+      }
+    }
+
+    class ThreadCountFieldMapper implements RowMapper<List<String>> {
       @Override
       public List<String> map(ResultSet rs, StatementContext ctx) throws SQLException {
-        return Arrays.asList(rs.getString("entityLink"), rs.getString("count"));
+        return Arrays.asList(
+            rs.getString("entityLink"),
+            rs.getString("type"),
+            rs.getString("taskStatus"),
+            rs.getString("count"));
       }
     }
   }

openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/FeedRepository.java

@@ -42,7 +42,6 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.UUID;
-import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 import javax.json.JsonPatch;
 import javax.ws.rs.core.Response.Status;
@@ -57,7 +56,6 @@
 import org.json.JSONObject;
 import org.openmetadata.schema.EntityInterface;
 import org.openmetadata.schema.api.feed.CloseTask;
-import org.openmetadata.schema.api.feed.EntityLinkThreadCount;
 import org.openmetadata.schema.api.feed.ResolveTask;
 import org.openmetadata.schema.api.feed.ThreadCount;
 import org.openmetadata.schema.entity.feed.Thread;
@@ -518,57 +516,86 @@ public void deleteByAbout(UUID entityId) {
     }
   }
 
-  public ThreadCount getThreadsCount(FeedFilter filter, String link) {
+  public List<ThreadCount> getThreadsCount(String link) {
     List<List<String>> result;
-    if (link == null) {
-      // Get thread count of all entities
+    EntityLink entityLink = EntityLink.parse(link);
+    List<ThreadCount> threadCounts = new ArrayList<>();
+    EntityReference reference = EntityUtil.validateEntityLink(entityLink);
+    int mentions;
+    if (reference.getType().equals(USER) || reference.getType().equals(Entity.TEAM)) {
+      if (reference.getType().equals(USER)) {
+        UUID userId = reference.getId();
+        User user = Entity.getEntity(USER, userId, TEAMS_FIELD, NON_DELETED);
+        List<String> teamIds = getTeamIds(user);
+        List<String> teamNames = getTeamNames(user);
+        result = dao.feedDAO().listCountByOwner(userId, teamIds);
+        mentions =
+            dao.feedDAO()
+                .listCountThreadsByMentions(
+                    FullyQualifiedName.buildHash(user.getFullyQualifiedName()),
+                    teamNames,
+                    Relationship.MENTIONED_IN.ordinal(),
+                    " where true ");
+      } else {
+        mentions = 0;
+        // team is not supported
+        result = new ArrayList<>();
+      }
+      ThreadCount threadCount = new ThreadCount().withMentionCount(mentions);
+      threadCount.setEntityLink(link);
+      result.forEach(
+          l -> {
+            String type = l.get(0);
+            String taskStatus = l.get(1);
+            int count = Integer.parseInt(l.get(2));
+            if (type.equalsIgnoreCase("Conversation")) {
+              threadCount.setConversationCount(count);
+            } else if (type.equalsIgnoreCase("Task")) {
+              if (taskStatus.equals("Open")) {
+                threadCount.setOpenTaskCount(count);
+              } else if (taskStatus.equals("Closed")) {
+                threadCount.setClosedTaskCount(count);
+              }
+            }
+          });
+      computeTotalTaskCount(threadCount);
+      threadCounts.add(threadCount);
+    } else {
+      mentions = 0;
       result =
-          // TODO fix this
           dao.feedDAO()
               .listCountByEntityLink(
-                  null,
-                  Entity.THREAD,
-                  null,
-                  IS_ABOUT.ordinal(),
-                  filter.getThreadType(),
-                  filter.getTaskStatus(),
-                  filter.getResolved());
-    } else {
-      EntityLink entityLink = EntityLink.parse(link);
-      EntityReference reference = EntityUtil.validateEntityLink(entityLink);
-      if (reference.getType().equals(USER) || reference.getType().equals(Entity.TEAM)) {
-        if (reference.getType().equals(USER)) {
-          UUID userId = reference.getId();
-          List<String> teamIds = getTeamIds(userId);
-          result = dao.feedDAO().listCountByOwner(userId, teamIds, filter.getCondition());
-        } else {
-          // team is not supported
-          result = new ArrayList<>();
-        }
-      } else {
-        result =
-            dao.feedDAO()
-                .listCountByEntityLink(
-                    entityLink.getFullyQualifiedFieldValue(),
-                    Entity.THREAD,
-                    entityLink.getFullyQualifiedFieldType(),
-                    IS_ABOUT.ordinal(),
-                    filter.getThreadType(),
-                    filter.getTaskStatus(),
-                    filter.getResolved());
-      }
+                  reference.getId(),
+                  reference.getFullyQualifiedName(),
+                  entityLink.getFullyQualifiedFieldType());
+      result.forEach(
+          l -> {
+            ThreadCount threadCount = new ThreadCount().withMentionCount(mentions);
+            String eLink = l.get(0);
+            String type = l.get(1);
+            String taskStatus = l.get(2);
+            threadCount.setEntityLink(eLink);
+            int count = Integer.parseInt(l.get(3));
+            if (type.equalsIgnoreCase("Conversation")) {
+              threadCount.setConversationCount(count);
+            } else if (type.equalsIgnoreCase("Task")) {
+              if (taskStatus.equals("Open")) {
+                threadCount.setOpenTaskCount(count);
+              } else if (taskStatus.equals("Closed")) {
+                threadCount.setClosedTaskCount(count);
+              }
+            }
+            computeTotalTaskCount(threadCount);
+            threadCounts.add(threadCount);
+          });
     }
+    return threadCounts;
+  }
 
-    AtomicInteger totalCount = new AtomicInteger(0);
-    List<EntityLinkThreadCount> entityLinkThreadCounts = new ArrayList<>();
-    result.forEach(
-        l -> {
-          int count = Integer.parseInt(l.get(1));
-          entityLinkThreadCounts.add(
-              new EntityLinkThreadCount().withEntityLink(l.get(0)).withCount(count));
-          totalCount.addAndGet(count);
-        });
-    return new ThreadCount().withTotalCount(totalCount.get()).withCounts(entityLinkThreadCounts);
+  private void computeTotalTaskCount(ThreadCount threadCount) {
+    threadCount.setTotalTaskCount(
+        (threadCount.getOpenTaskCount() != null ? threadCount.getOpenTaskCount() : 0)
+            + (threadCount.getClosedTaskCount() != null ? threadCount.getClosedTaskCount() : 0));
   }
 
   public List<Post> listPosts(UUID threadId) {
@@ -1084,11 +1111,17 @@ private List<String> getTeamIds(UUID userId) {
     List<String> teamIds = null;
     if (userId != null) {
       User user = Entity.getEntity(Entity.USER, userId, TEAMS_FIELD, NON_DELETED);
-      teamIds = listOrEmpty(user.getTeams()).stream().map(ref -> ref.getId().toString()).toList();
+      teamIds = getTeamIds(user);
     }
     return nullOrEmpty(teamIds) ? List.of(StringUtils.EMPTY) : teamIds;
   }
 
+  private List<String> getTeamIds(User user) {
+    List<String> teamIds =
+        listOrEmpty(user.getTeams()).stream().map(ref -> ref.getId().toString()).toList();
+    return nullOrEmpty(teamIds) ? List.of(StringUtils.EMPTY) : teamIds;
+  }
+
   /** Returns the threads that are associated with the entities followed by the user. */
   private FilteredThreads getThreadsByFollows(FeedFilter filter, UUID userId, int limit) {
     List<String> teamIds = getTeamIds(userId);

openmetadata-service/src/main/java/org/openmetadata/service/resources/feeds/FeedResource.java

@@ -116,6 +116,10 @@ public static class PostList extends ResultList<Post> {
     /* Required for serde */
   }
 
+  public static class ThreadCountList extends ResultList<ThreadCount> {
+    /* Required for serde */
+  }
+
   @GET
   @Operation(
       operationId = "listThreads",
@@ -373,41 +377,17 @@ public Response updateThread(
             content =
                 @Content(
                     mediaType = "application/json",
-                    schema = @Schema(implementation = ThreadCount.class)))
+                    schema = @Schema(implementation = ThreadCountList.class)))
       })
-  public ThreadCount getThreadCount(
+  public ResultList<ThreadCount> getThreadCount(
       @Context UriInfo uriInfo,
       @Parameter(
               description = "Filter threads by entity link",
               schema =
                   @Schema(type = "string", example = "<E#/{entityType}/{entityFQN}/{fieldName}>"))
           @QueryParam("entityLink")
-          String entityLink,
-      @Parameter(
-              description =
-                  "The type of thread to filter the results. It can take one of 'Conversation', 'Task', 'Announcement'",
-              schema = @Schema(implementation = ThreadType.class))
-          @QueryParam("type")
-          ThreadType threadType,
-      @Parameter(
-              description =
-                  "The status of tasks to filter the results. It can take one of 'Open', 'Closed'. This filter will take effect only when type is set to Task",
-              schema = @Schema(implementation = TaskStatus.class))
-          @QueryParam("taskStatus")
-          TaskStatus taskStatus,
-      @Parameter(
-              description = "Filter threads by whether it is active or resolved",
-              schema = @Schema(type = "boolean"))
-          @DefaultValue("false")
-          @QueryParam("isResolved")
-          Boolean isResolved) {
-    FeedFilter filter =
-        FeedFilter.builder()
-            .threadType(threadType)
-            .taskStatus(taskStatus)
-            .resolved(isResolved)
-            .build();
-    return dao.getThreadsCount(filter, entityLink);
+          String entityLink) {
+    return new ResultList<>(dao.getThreadsCount(entityLink));
   }
 
   @POST

openmetadata-service/src/test/java/org/openmetadata/service/resources/feeds/FeedResourceTest.java

@@ -79,7 +79,6 @@
 import org.openmetadata.schema.api.feed.CloseTask;
 import org.openmetadata.schema.api.feed.CreatePost;
 import org.openmetadata.schema.api.feed.CreateThread;
-import org.openmetadata.schema.api.feed.EntityLinkThreadCount;
 import org.openmetadata.schema.api.feed.ResolveTask;
 import org.openmetadata.schema.api.feed.ThreadCount;
 import org.openmetadata.schema.api.teams.CreateTeam;
@@ -348,11 +347,16 @@
     // Test the /api/v1/feed/count API
     assertEquals(
         userThreadCount, listThreads(USER_LINK, null, USER_AUTH_HEADERS).getPaging().getTotal());
-    assertEquals(userThreadCount, listThreadsCount(USER_LINK, USER_AUTH_HEADERS).getTotalCount());
+    FeedResource.ThreadCountList threadCounts = listThreadsCount(USER_LINK, USER_AUTH_HEADERS);
+    for (ThreadCount threadCount : threadCounts.getData()) {
+      if (threadCount.getEntityLink().equals(USER_LINK)) {
+        assertEquals(userThreadCount, threadCount.getConversationCount());
+      }
+    }
     assertEquals(
         tableDescriptionThreadCount, getThreadCount(TABLE_DESCRIPTION_LINK, USER_AUTH_HEADERS));
     assertEquals(
         tableColumnDescriptionThreadCount, getThreadCount(TABLE_COLUMN_LINK, USER_AUTH_HEADERS));
   }
 
   @Test
@@ -441,23 +445,15 @@
     assertEquals(assignedToCount + assignedByCount + 2, tasks.getPaging().getTotal());
     assertEquals(assignedToCount + assignedByCount + 2, tasks.getData().size());
 
-    ThreadCount count = listTasksCount(null, TaskStatus.Open, USER2_AUTH_HEADERS);
-    int totalOpenTaskCount = count.getTotalCount();
-    count = listTasksCount(null, TaskStatus.Closed, USER2_AUTH_HEADERS);
-    int totalClosedTaskCount = count.getTotalCount();
-
     // close a task and test the task status filter
     ResolveTask resolveTask = new ResolveTask().withNewValue("accepted description");
     resolveTask(task2.getId(), resolveTask, USER_AUTH_HEADERS);
 
     tasks = listTasks(null, null, null, TaskStatus.Open, null, USER2_AUTH_HEADERS);
     assertFalse(tasks.getData().stream().anyMatch(t -> t.getTask().getId().equals(task2.getId())));
-    assertEquals(totalOpenTaskCount - 1, tasks.getPaging().getTotal());
 
     tasks = listTasks(null, null, null, TaskStatus.Closed, null, USER2_AUTH_HEADERS);
     assertEquals(task2.getId(), tasks.getData().get(0).getTask().getId());
-    assertEquals(totalClosedTaskCount + 1, tasks.getPaging().getTotal());
-    assertEquals(totalClosedTaskCount + 1, tasks.getData().size());
   }
 
   @Test
@@ -1622,34 +1618,30 @@
     return TestUtils.get(target, PostList.class, authHeaders);
   }
 
-  public ThreadCount listThreadsCount(String entityLink, Map<String, String> authHeaders)
-      throws HttpResponseException {
+  public FeedResource.ThreadCountList listThreadsCount(
+      String entityLink, Map<String, String> authHeaders) throws HttpResponseException {
     WebTarget target = getResource("feed/count");
     target = entityLink != null ? target.queryParam("entityLink", entityLink) : target;
-    target = target.queryParam("type", ThreadType.Conversation);
-    return TestUtils.get(target, ThreadCount.class, authHeaders);
+    return TestUtils.get(target, FeedResource.ThreadCountList.class, authHeaders);
   }
 
   public ThreadCount listTasksCount(
       String entityLink, TaskStatus taskStatus, Map<String, String> authHeaders)
       throws HttpResponseException {
     WebTarget target = getResource("feed/count");
     target = entityLink != null ? target.queryParam("entityLink", entityLink) : target;
-    target = target.queryParam("type", ThreadType.Task);
-    target = taskStatus != null ? target.queryParam("taskStatus", taskStatus) : target;
     return TestUtils.get(target, ThreadCount.class, authHeaders);
   }
 
   private int getThreadCount(String entityLink, Map<String, String> authHeaders)
       throws HttpResponseException {
-    List<EntityLinkThreadCount> linkThreadCount =
-        listThreadsCount(entityLink, authHeaders).getCounts();
-    EntityLinkThreadCount threadCount =
-        linkThreadCount.stream()
-            .filter(l -> l.getEntityLink().equals(entityLink))
-            .findFirst()
-            .orElseThrow();
-    return threadCount.getCount();
+    FeedResource.ThreadCountList threadCounts = listThreadsCount(entityLink, authHeaders);
+    for (ThreadCount threadCount : threadCounts.getData()) {
+      if (threadCount.getEntityLink().equalsIgnoreCase(entityLink)) {
+        return threadCount.getConversationCount();
+      }
+    }
+    return 0;
   }
 
   protected final Thread patchThreadAndCheck(