Fix #22422: Classification term count showing 0 in UI … (#22470)

* Fix #22422: Classification term count showing 0 in UI                                                                                                                   │
│                                                                                                                                                                                          │
│   - Added bulk fetch methods to ClassificationDAO for efficient term count retrieval                                                                                                     │
│   - Implemented proper batch fetching in ClassificationRepository.setFieldsInBulk                                                                                                        │
│   - Added ConnectionAwareSqlQuery annotations for MySQL and PostgreSQL compatibility                                                                                                     │
│   - Fixed duplicate key error in TagRepository.batchFetchUsageCounts                                                                                                                     │
│   - Added Playwright test to verify classification term counts display correctly

* Fix java checkstyle

* remove unused query

* refactor query :  switch to join-based classification hash matching and add generated column for classificationHash

---------

Co-authored-by: sonikashah <sonikashah94@gmail.com>
Co-authored-by: sonika-shah <58761340+sonika-shah@users.noreply.github.com>

bootstrap/sql/migrations/native/1.9.0/mysql/schemaChanges.sql

@@ -136,3 +136,13 @@ CREATE TABLE IF NOT EXISTS entity_deletion_lock (
     INDEX idx_deletion_lock_fqn (entityFqn(255)),
     INDEX idx_deletion_lock_time (lockedAt)
 ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci;
+
+
+-- 1. Add generated classificationHash column to support fast lookup and grouping by classification fqnHash
+ALTER TABLE tag
+  ADD COLUMN classificationHash VARCHAR(255)
+  GENERATED ALWAYS AS (SUBSTRING_INDEX(fqnhash, '.', 1)) STORED;
+
+-- 2. Create index on classificationHash + deleted
+CREATE INDEX idx_tag_classification_hash_deleted
+  ON tag (classificationHash, deleted);

bootstrap/sql/migrations/native/1.9.0/postgres/schemaChanges.sql

@@ -170,3 +170,12 @@ CREATE TABLE IF NOT EXISTS entity_deletion_lock (
 -- Use btree index for entityFqn prefix matching
 CREATE INDEX IF NOT EXISTS idx_deletion_lock_fqn ON entity_deletion_lock(entityFqn);
 CREATE INDEX IF NOT EXISTS idx_deletion_lock_time ON entity_deletion_lock(lockedAt);
+
+-- 1. Add classificationHash column to support fast lookup and grouping by classification fqnHash
+ALTER TABLE tag
+  ADD COLUMN classificationHash TEXT
+  GENERATED ALWAYS AS (SPLIT_PART(fqnhash, '.', 1)) STORED;
+
+-- 2. Create index on classificationHash + deleted
+CREATE INDEX idx_tag_classification_hash_deleted
+  ON tag (classificationHash, deleted);

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

@@ -19,11 +19,16 @@ import static org.openmetadata.service.search.SearchClient.TAG_SEARCH_INDEX;
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.apache.commons.lang3.tuple.Pair;
 import org.jdbi.v3.core.mapper.RowMapper;
 import org.jdbi.v3.sqlobject.transaction.Transaction;
 import org.openmetadata.schema.entity.classification.Classification;
@@ -41,6 +46,7 @@ import org.openmetadata.service.exception.CatalogExceptionMessage;
 import org.openmetadata.service.jdbi3.CollectionDAO.EntityRelationshipRecord;
 import org.openmetadata.service.resources.tags.ClassificationResource;
 import org.openmetadata.service.util.EntityUtil.Fields;
+import org.openmetadata.service.util.FullyQualifiedName;
 
 @Slf4j
 public class ClassificationRepository extends EntityRepository<Classification> {
@@ -82,6 +88,124 @@ public class ClassificationRepository extends EntityRepository<Classification> {
         fields.contains("usageCount") ? classification.getUsageCount() : null);
   }
 
+  @Override
+  public void setFieldsInBulk(Fields fields, List<Classification> entities) {
+    if (entities == null || entities.isEmpty()) {
+      return;
+    }
+    fetchAndSetFields(entities, fields);
+    fetchAndSetClassificationSpecificFields(entities, fields);
+    setInheritedFields(entities, fields);
+    for (Classification entity : entities) {
+      clearFieldsInternal(entity, fields);
+    }
+  }
+
+  private void fetchAndSetClassificationSpecificFields(
+      List<Classification> classifications, Fields fields) {
+    if (classifications == null || classifications.isEmpty()) {
+      return;
+    }
+    if (fields.contains("termCount")) {
+      fetchAndSetTermCounts(classifications);
+    }
+    if (fields.contains("usageCount")) {
+      fetchAndSetUsageCounts(classifications);
+    }
+  }
+
+  private void fetchAndSetTermCounts(List<Classification> classifications) {
+    // Batch fetch term counts for all classifications
+    Map<String, Integer> termCountMap = batchFetchTermCounts(classifications);
+    for (Classification classification : classifications) {
+      classification.withTermCount(
+          termCountMap.getOrDefault(classification.getFullyQualifiedName(), 0));
+    }
+  }
+
+  private void fetchAndSetUsageCounts(List<Classification> classifications) {
+    Map<String, Integer> usageCountMap = batchFetchUsageCounts(classifications);
+    for (Classification classification : classifications) {
+      classification.withUsageCount(
+          usageCountMap.getOrDefault(classification.getFullyQualifiedName(), 0));
+    }
+  }
+
+  private Map<String, Integer> batchFetchTermCounts(List<Classification> classifications) {
+    Map<String, Integer> termCountMap = new HashMap<>();
+    if (classifications == null || classifications.isEmpty()) {
+      return termCountMap;
+    }
+
+    try {
+      // Convert classifications to their hash representations
+      List<String> classificationHashes = new ArrayList<>();
+      Map<String, String> hashToFqnMap = new HashMap<>();
+
+      for (Classification classification : classifications) {
+        String fqn = classification.getFullyQualifiedName();
+        String hash = FullyQualifiedName.buildHash(fqn);
+        classificationHashes.add(hash);
+        hashToFqnMap.put(hash, fqn);
+      }
+
+      // Convert the list to JSON array for MySQL
+      String jsonHashes = JsonUtils.pojoToJson(classificationHashes);
+
+      // Use the DAO method for database-specific query
+      List<Pair<String, Integer>> results =
+          daoCollection.classificationDAO().bulkGetTermCounts(jsonHashes, classificationHashes);
+
+      // Process results
+      for (Pair<String, Integer> result : results) {
+        String classificationHash = result.getLeft();
+        Integer count = result.getRight();
+        String fqn = hashToFqnMap.get(classificationHash);
+        if (fqn != null) {
+          termCountMap.put(fqn, count);
+        }
+      }
+
+      // Set 0 for classifications with no tags
+      for (Classification classification : classifications) {
+        termCountMap.putIfAbsent(classification.getFullyQualifiedName(), 0);
+      }
+
+      return termCountMap;
+    } catch (Exception e) {
+      LOG.error("Error batch fetching term counts, falling back to individual queries", e);
+      // Fall back to individual queries
+      for (Classification classification : classifications) {
+        ListFilter filterWithParent =
+            new ListFilter(Include.NON_DELETED)
+                .addQueryParam("parent", classification.getFullyQualifiedName());
+        int count = daoCollection.tagDAO().listCount(filterWithParent);
+        termCountMap.put(classification.getFullyQualifiedName(), count);
+      }
+      return termCountMap;
+    }
+  }
+
+  private Map<String, Integer> batchFetchUsageCounts(List<Classification> classifications) {
+    Map<String, Integer> usageCountMap = new HashMap<>();
+    if (classifications == null || classifications.isEmpty()) {
+      return usageCountMap;
+    }
+
+    // Batch fetch usage counts for all classifications at once
+    List<String> classificationFQNs =
+        classifications.stream()
+            .map(Classification::getFullyQualifiedName)
+            .collect(Collectors.toList());
+
+    Map<String, Integer> counts =
+        daoCollection
+            .tagUsageDAO()
+            .getTagCountsBulk(TagSource.CLASSIFICATION.ordinal(), classificationFQNs);
+
+    return counts != null ? counts : usageCountMap;
+  }
+
   @Override
   public void prepare(Classification entity, boolean update) {
     /* Nothing to do */

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

@@ -3688,6 +3688,40 @@ public interface CollectionDAO {
     default String getNameHashColumn() {
       return "nameHash";
     }
+
+    @ConnectionAwareSqlQuery(
+        value =
+            ""
+                + "SELECT t.classificationHash, COUNT(*) AS termCount "
+                + "FROM tag t "
+                + "JOIN JSON_TABLE("
+                + "    :jsonHashes, "
+                + "    '$[*]' COLUMNS (classificationHash VARCHAR(64) PATH '$')"
+                + ") AS h "
+                + "  ON t.classificationHash = h.classificationHash "
+                + "WHERE t.deleted = FALSE "
+                + "GROUP BY t.classificationHash",
+        connectionType = MYSQL)
+    @ConnectionAwareSqlQuery(
+        value =
+            ""
+                + "SELECT t.classificationHash, COUNT(*) AS termCount "
+                + "FROM tag t "
+                + "JOIN UNNEST(:hashArray::text[]) AS h(classificationHash) "
+                + "  ON t.classificationHash = h.classificationHash "
+                + "WHERE t.deleted = FALSE "
+                + "GROUP BY t.classificationHash",
+        connectionType = POSTGRES)
+    @RegisterRowMapper(TermCountMapper.class)
+    List<Pair<String, Integer>> bulkGetTermCounts(
+        @Bind("jsonHashes") String jsonHashes, @Bind("hashArray") List<String> hashArray);
+
+    class TermCountMapper implements RowMapper<Pair<String, Integer>> {
+      @Override
+      public Pair<String, Integer> map(ResultSet rs, StatementContext ctx) throws SQLException {
+        return Pair.of(rs.getString("classificationHash"), rs.getInt("termCount"));
+      }
+    }
   }
 
   interface TagDAO extends EntityDAO<Tag> {
@@ -3708,6 +3742,23 @@ public interface CollectionDAO {
 
     @Override
     default int listCount(ListFilter filter) {
+      String parent = filter.getQueryParam("parent");
+
+      // If parent parameter is provided, filter tags by parent classification FQN
+      if (parent != null) {
+        String parentFqnHash = FullyQualifiedName.buildHash(parent);
+        filter.queryParams.put("parentFqnPrefix", parentFqnHash + ".%");
+        String condition = filter.getCondition("tag");
+        if (!condition.isEmpty()) {
+          condition = String.format("%s AND fqnHash LIKE :parentFqnPrefix", condition);
+        } else {
+          condition = "WHERE fqnHash LIKE :parentFqnPrefix";
+        }
+        return listCount(
+            getTableName(), getNameHashColumn(), filter.getQueryParams(), condition, condition);
+      }
+
+      // Original behavior for classification.disabled parameter
       boolean disabled = Boolean.parseBoolean(filter.getQueryParam("classification.disabled"));
       String condition =
           String.format(
@@ -3750,6 +3801,29 @@ public interface CollectionDAO {
     @Override
     default List<String> listBefore(
         ListFilter filter, int limit, String beforeName, String beforeId) {
+      String parent = filter.getQueryParam("parent");
+
+      // If parent parameter is provided, filter tags by parent classification FQN
+      if (parent != null) {
+        String parentFqnHash = FullyQualifiedName.buildHash(parent);
+        filter.queryParams.put("parentFqnPrefix", parentFqnHash + ".%");
+        String condition = filter.getCondition("tag");
+        if (!condition.isEmpty()) {
+          condition = String.format("%s AND fqnHash LIKE :parentFqnPrefix", condition);
+        } else {
+          condition = "WHERE fqnHash LIKE :parentFqnPrefix";
+        }
+        return listBefore(
+            getTableName(),
+            filter.getQueryParams(),
+            condition,
+            condition,
+            limit,
+            beforeName,
+            beforeId);
+      }
+
+      // Original behavior for classification.disabled parameter
       boolean disabled = Boolean.parseBoolean(filter.getQueryParam("classification.disabled"));
       String condition =
           String.format(
@@ -3795,6 +3869,29 @@ public interface CollectionDAO {
 
     @Override
     default List<String> listAfter(ListFilter filter, int limit, String afterName, String afterId) {
+      String parent = filter.getQueryParam("parent");
+
+      // If parent parameter is provided, filter tags by parent classification FQN
+      if (parent != null) {
+        String parentFqnHash = FullyQualifiedName.buildHash(parent);
+        filter.queryParams.put("parentFqnPrefix", parentFqnHash + ".%");
+        String condition = filter.getCondition("tag");
+        if (!condition.isEmpty()) {
+          condition = String.format("%s AND fqnHash LIKE :parentFqnPrefix", condition);
+        } else {
+          condition = "WHERE fqnHash LIKE :parentFqnPrefix";
+        }
+        return listAfter(
+            getTableName(),
+            filter.getQueryParams(),
+            condition,
+            condition,
+            limit,
+            afterName,
+            afterId);
+      }
+
+      // Original behavior for classification.disabled parameter
       boolean disabled = Boolean.parseBoolean(filter.getQueryParam("classification.disabled"));
       String condition =
           String.format(
@@ -3910,16 +4007,14 @@ public interface CollectionDAO {
         value =
             "SELECT source, tagFQN, labelType, targetFQNHash, state, json "
                 + "FROM ("
-                + "  SELECT gterm.* , tu.* "
+                + "  SELECT tu.source, tu.tagFQN, tu.labelType, tu.targetFQNHash, tu.state, gterm.json "
                 + "  FROM glossary_term_entity AS gterm "
-                + "  JOIN tag_usage AS tu "
+                + "  JOIN tag_usage AS tu ON gterm.fqnHash = tu.tagFQNHash "
-                + "  ON gterm.fqnHash = tu.tagFQNHash "
                 + "  WHERE tu.source = 1 "
                 + "  UNION ALL "
-                + "  SELECT ta.*, tu.* "
+                + "  SELECT tu.source, tu.tagFQN, tu.labelType, tu.targetFQNHash, tu.state, ta.json "
                 + "  FROM tag AS ta "
-                + "  JOIN tag_usage AS tu "
+                + "  JOIN tag_usage AS tu ON ta.fqnHash = tu.tagFQNHash "
-                + "  ON ta.fqnHash = tu.tagFQNHash "
                 + "  WHERE tu.source = 0 "
                 + ") AS combined_data "
                 + "WHERE combined_data.targetFQNHash LIKE :targetFQNHash",
@@ -3928,12 +4023,12 @@ public interface CollectionDAO {
         value =
             "SELECT source, tagFQN, labelType, targetFQNHash, state, json "
                 + "FROM ("
-                + "  SELECT gterm.*, tu.* "
+                + "  SELECT tu.source, tu.tagFQN, tu.labelType, tu.targetFQNHash, tu.state, gterm.json "
                 + "  FROM glossary_term_entity AS gterm "
                 + "  JOIN tag_usage AS tu ON gterm.fqnHash = tu.tagFQNHash "
                 + "  WHERE tu.source = 1 "
                 + "  UNION ALL "
-                + "  SELECT ta.*, tu.* "
+                + "  SELECT tu.source, tu.tagFQN, tu.labelType, tu.targetFQNHash, tu.state, ta.json "
                 + "  FROM tag AS ta "
                 + "  JOIN tag_usage AS tu ON ta.fqnHash = tu.tagFQNHash "
                 + "  WHERE tu.source = 0 "

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

@@ -431,6 +431,7 @@ public class TagRepository extends EntityRepository<Tag> {
               .withHandle(handle -> handle.createQuery(queryBuilder.toString()).mapToMap().list());
 
       return results.stream()
+          .filter(row -> row.get("tagFQN") != null)
           .collect(
               Collectors.toMap(
                   row -> (String) row.get("tagFQN"),

openmetadata-service/src/test/java/org/openmetadata/service/resources/tags/ClassificationResourceTest.java

@@ -217,4 +217,26 @@ public class ClassificationResourceTest
       assertTrue(child.getFullyQualifiedName().startsWith(ret.getFullyQualifiedName()));
     }
   }
+
+  @Test
+  void testClassificationTermCount(TestInfo test) throws IOException {
+    // Create a new classification
+    CreateClassification create = createRequest(getEntityName(test));
+    Classification classification = createAndCheckEntity(create, ADMIN_AUTH_HEADERS);
+
+    // Initially, termCount should be 0 when requested with termCount field
+    Classification withTermCount =
+        getEntity(classification.getId(), "termCount", ADMIN_AUTH_HEADERS);
+    assertEquals(0, withTermCount.getTermCount(), "New classification should have 0 tags");
+
+    // Create tags under this classification
+    TagResourceTest tagResourceTest = new TagResourceTest();
+    for (int i = 1; i <= 3; i++) {
+      tagResourceTest.createTag("tag" + i, classification.getName(), null);
+    }
+
+    // Now check termCount again
+    withTermCount = getEntity(classification.getId(), "termCount", ADMIN_AUTH_HEADERS);
+    assertEquals(3, withTermCount.getTermCount(), "Classification should have 3 tags");
+  }
 }

openmetadata-ui/src/main/resources/ui/playwright/e2e/Pages/Tags.spec.ts

@@ -289,6 +289,35 @@ test.fixme('Classification Page', async ({ page }) => {
     );
   });
 
+  await test.step('Verify classification term count', async () => {
+    // Navigate back to classifications list
+    await sidebarClick(page, SidebarItem.TAGS);
+
+    // Wait for classifications to load with termCount field
+    const classificationsResponse = page.waitForResponse(
+      (response) =>
+        response.url().includes('/api/v1/classifications') &&
+        response.url().includes('fields=termCount')
+    );
+    await classificationsResponse;
+
+    // Find the classification in the left panel and verify term count
+    const classificationElement = page
+      .locator(`[data-testid="side-panel-classification"]`)
+      .filter({ hasText: NEW_CLASSIFICATION.displayName });
+
+    // Check if term count is displayed as (1) since we created one tag
+    await expect(classificationElement).toContainText('(1)');
+
+    // Click on the classification to verify
+    await classificationElement.click();
+
+    // Verify the tag is listed
+    await expect(page.locator('[data-testid="table"]')).toContainText(
+      NEW_TAG.name
+    );
+  });
+
   await test.step(`Assign tag to table`, async () => {
     await table.visitEntityPage(page);
     const { name, displayName } = NEW_TAG;
@@ -423,6 +452,25 @@ test.fixme('Classification Page', async ({ page }) => {
     await expect(page.locator('[data-testid="table"]')).not.toContainText(
       NEW_TAG.name
     );
+
+    // Verify term count is now 0 after deleting the tag
+    await sidebarClick(page, SidebarItem.TAGS);
+
+    // Wait for classifications to reload with updated termCount
+    const classificationsResponse = page.waitForResponse(
+      (response) =>
+        response.url().includes('/api/v1/classifications') &&
+        response.url().includes('fields=termCount')
+    );
+    await classificationsResponse;
+
+    // Find the classification and verify term count is 0
+    const classificationElement = page
+      .locator(`[data-testid="side-panel-classification"]`)
+      .filter({ hasText: NEW_CLASSIFICATION.displayName });
+
+    // Check if term count is displayed as (0) since we deleted the tag
+    await expect(classificationElement).toContainText('(0)');
   });
 
   await test.step('Remove classification', async () => {
@@ -505,6 +553,57 @@ test('Search tag using classification display name should work', async ({
   ).toBeVisible();
 });
 
+test('Verify system classification term counts', async ({ page }) => {
+  // Navigate to tags page
+  await sidebarClick(page, SidebarItem.TAGS);
+
+  // Wait for classifications to load with termCount field
+  const classificationsResponse = page.waitForResponse(
+    (response) =>
+      response.url().includes('/api/v1/classifications') &&
+      response.url().includes('fields=termCount')
+  );
+  await classificationsResponse;
+
+  // Wait a bit for the UI to update
+  await page.waitForTimeout(1000);
+
+  // Get all classification elements
+  const classificationElements = await page
+    .locator('[data-testid="side-panel-classification"]')
+    .all();
+
+  // Find and verify Tier classification
+  let tierFound = false;
+  let piiFound = false;
+
+  for (const element of classificationElements) {
+    const text = await element.textContent();
+    if (text?.includes('Tier') && text?.includes('5')) {
+      tierFound = true;
+    }
+    if (text?.includes('PII') && text?.includes('3')) {
+      piiFound = true;
+    }
+  }
+
+  expect(tierFound).toBeTruthy();
+  expect(piiFound).toBeTruthy();
+
+  // Alternative: verify specific classifications have the expected count
+  const tierElement = page
+    .locator('[data-testid="side-panel-classification"]')
+    .filter({ hasText: 'Tier' });
+
+  await expect(tierElement).toContainText('5');
+
+  const piiElement = page
+    .locator('[data-testid="side-panel-classification"]')
+    .filter({ hasText: 'PII' });
+
+  await expect(piiElement).toContainText('3');
+});
+
 test('Verify Owner Add Delete', async ({ page }) => {
   await classification1.visitPage(page);
   const OWNER1 = user1.getUserName();