fix: [DHIS2-18459] Re-use obtained maxTeLimit (2.39)

dhis-2/dhis-api/src/main/java/org/hisp/dhis/trackedentity/TrackedEntityInstanceQueryParams.java

@@ -504,6 +504,11 @@ public boolean hasTrackedEntityType() {
     return trackedEntityType != null;
   }
 
+  /** Indicates whether this parameters specifies a max TE limit. */
+  public boolean hasMaxTeiLimit() {
+    return maxTeiLimit > 0;
+  }
+
   /** Indicates whether this parameters is of the given organisation unit mode. */
   public boolean isOrganisationUnitMode(OrganisationUnitSelectionMode mode) {
     return organisationUnitMode != null && organisationUnitMode.equals(mode);

dhis-2/dhis-api/src/main/java/org/hisp/dhis/trackedentity/TrackedEntityInstanceStore.java

@@ -49,6 +49,8 @@ public interface TrackedEntityInstanceStore extends IdentifiableObjectStore<Trac
 
   int getTrackedEntityInstanceCountForGrid(TrackedEntityInstanceQueryParams params);
 
+  int getTrackedEntityInstanceCountForGridWithMaxTeiLimit(TrackedEntityInstanceQueryParams params);
+
   /**
    * Checks for the existence of a TEI by UID. Deleted TEIs are not taken into account.
    *

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/trackedentity/DefaultTrackedEntityInstanceService.java

@@ -739,16 +739,19 @@ public void validateSearchScope(TrackedEntityInstanceQueryParams params, boolean
         }
       }
 
-      checkIfMaxTeiLimitIsReached(params, maxTeiLimit);
       params.setMaxTeiLimit(maxTeiLimit);
+      checkIfMaxTeiLimitIsReached(params);
     }
   }
 
-  private void checkIfMaxTeiLimitIsReached(
-      TrackedEntityInstanceQueryParams params, int maxTeiLimit) {
-    int instanceCount = trackedEntityInstanceStore.getTrackedEntityInstanceCountForGrid(params);
-    if (maxTeiLimit > 0 && instanceCount > maxTeiLimit) {
-      throw new IllegalQueryException("maxteicountreached");
+  private void checkIfMaxTeiLimitIsReached(TrackedEntityInstanceQueryParams params) {
+    if (params.hasMaxTeiLimit()) {
+      int instanceCount =
+          trackedEntityInstanceStore.getTrackedEntityInstanceCountForGridWithMaxTeiLimit(params);
+
+      if (instanceCount > params.getMaxTeiLimit()) {
+        throw new IllegalQueryException("maxteicountreached");
+      }
     }
   }
 

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/trackedentity/hibernate/HibernateTrackedEntityInstanceStore.java

@@ -322,6 +322,24 @@ public int getTrackedEntityInstanceCountForGrid(TrackedEntityInstanceQueryParams
     return jdbcTemplate.queryForObject(sql, Integer.class);
   }
 
+  @Override
+  public int getTrackedEntityInstanceCountForGridWithMaxTeiLimit(
+      TrackedEntityInstanceQueryParams params) {
+    // ---------------------------------------------------------------------
+    // Select clause
+    // ---------------------------------------------------------------------
+
+    String sql = getCountQueryWithMaxTeiLimit(params);
+
+    // ---------------------------------------------------------------------
+    // Query
+    // ---------------------------------------------------------------------
+
+    log.debug("Tracked entity instance count SQL: " + sql);
+
+    return jdbcTemplate.queryForObject(sql, Integer.class);
+  }
+
   /**
    * Generates SQL based on "params". The purpose of the SQL is to retrieve a list of tracked entity
    * instances, and additionally any requested attributes (If defined in params).
@@ -416,6 +434,26 @@ private String getCountQuery(TrackedEntityInstanceQueryParams params) {
         .toString();
   }
 
+  /**
+   * Uses the same basis as the getQuery method, but replaces the projection with a count, ignores
+   * order but uses the TEI limit set on the program if higher than 0
+   *
+   * @param params params defining the query
+   * @return a count SQL query
+   */
+  private String getCountQueryWithMaxTeiLimit(TrackedEntityInstanceQueryParams params) {
+    return new StringBuilder()
+        .append(getQueryCountSelect(params))
+        .append(getQuerySelect(params))
+        .append("FROM ")
+        .append(getFromSubQuery(params, true, true))
+        .append(getQueryRelatedTables(params))
+        .append(getQueryGroupBy(params))
+        .append(params.hasMaxTeiLimit() ? getLimitClause(params.getMaxTeiLimit() + 1) : "")
+        .append(" ) teicount")
+        .toString();
+  }
+
   /**
    * Generates the projection of the main query. Includes two optional columns, deleted and
    * tea_values
@@ -1209,6 +1247,10 @@ private String getQueryGroupBy(TrackedEntityInstanceQueryParams params) {
     return groupBy.toString();
   }
 
+  private String getLimitClause(int limit) {
+    return "LIMIT " + limit;
+  }
+
   /**
    * Generates the ORDER BY clause. This clause is used both in the subquery and main query. When
    * using it in the subquery, we want to make sure we get the right teis. When we order in the main

dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/trackedentity/TrackedEntityInstanceStoreTest.java

@@ -40,6 +40,7 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import org.hisp.dhis.analytics.AggregationType;
 import org.hisp.dhis.common.OrganisationUnitSelectionMode;
 import org.hisp.dhis.common.QueryItem;
@@ -121,6 +122,7 @@ public void setUpTest() {
     organisationUnitService.addOrganisationUnit(ouB);
     organisationUnitService.addOrganisationUnit(ouC);
     trackedEntityType = createTrackedEntityType('A');
+    trackedEntityType.setMaxTeiCountToReturn(1000);
     trackedEntityTypeService.addTrackedEntityType(trackedEntityType);
     prA = createProgram('A', null, null);
     prA.setTrackedEntityType(trackedEntityType);
@@ -519,6 +521,32 @@ void testPotentialDuplicateInGridQuery() {
         });
   }
 
+  @Test
+  void shouldHandleNullProgramWhenCheckingMaxTeiLimitInGridQuery() {
+    trackedEntityTypeService.addTrackedEntityType(trackedEntityType);
+    teiA.setTrackedEntityType(trackedEntityType);
+    teiA.setPotentialDuplicate(true);
+    teiStore.save(teiA);
+    teiB.setTrackedEntityType(trackedEntityType);
+    teiB.setPotentialDuplicate(true);
+    teiStore.save(teiB);
+    teiC.setTrackedEntityType(trackedEntityType);
+    teiStore.save(teiC);
+    teiD.setTrackedEntityType(trackedEntityType);
+    teiStore.save(teiD);
+    dbmsManager.flushSession();
+
+    TrackedEntityInstanceQueryParams params = new TrackedEntityInstanceQueryParams();
+    params.setOrganisationUnitMode(OrganisationUnitSelectionMode.ACCESSIBLE);
+    params.setPrograms(List.of(prA, prB));
+    params.setTrackedEntityType(trackedEntityType);
+    params.setTrackedEntityInstanceUids(
+        Set.of(teiA.getUid(), teiB.getUid(), teiC.getUid(), teiD.getUid()));
+
+    int count = teiStore.getTrackedEntityInstanceCountForGridWithMaxTeiLimit(params);
+    assertEquals(4, count);
+  }
+
   @Test
   void testProgramAttributeOfTypeOrgUnitIsResolvedToOrgUnitName() {
     trackedEntityTypeService.addTrackedEntityType(trackedEntityType);