perf: DataEntry metadata with fewer n+1s [DHIS2-21757] (#24509)

Author: david-mackessy

dhis-2/dhis-api/src/main/java/org/hisp/dhis/category/CategoryOptionComboStore.java

@@ -44,6 +44,18 @@ public interface CategoryOptionComboStore extends IdentifiableObjectStore<Catego
   CategoryOptionCombo getCategoryOptionCombo(
       CategoryCombo categoryCombo, Set<CategoryOption> categoryOptions);
 
+  /**
+   * Returns the {@link CategoryOptionCombo}s of the given {@link CategoryCombo}s with their {@link
+   * CategoryOptionCombo#getCategoryOptions()} eagerly fetched in a single query. Used to prime the
+   * session before serialising category option combos, avoiding the per-combo N+1 select on the
+   * {@code categoryoptioncombos_categoryoptions} join table.
+   *
+   * @param categoryCombos the category combos whose option combos to load.
+   * @return the option combos with their category options initialised.
+   */
+  List<CategoryOptionCombo> getCategoryOptionCombosWithCategoryOptions(
+      Collection<CategoryCombo> categoryCombos);
+
   void updateNames();
 
   void deleteNoRollBack(CategoryOptionCombo categoryOptionCombo);

dhis-2/dhis-api/src/main/java/org/hisp/dhis/category/CategoryService.java

@@ -412,6 +412,18 @@ CategoryOptionCombo getCategoryOptionCombo(
    */
   List<CategoryOptionCombo> getAllCategoryOptionCombos();
 
+  /**
+   * Returns the {@link CategoryOptionCombo}s of the given {@link CategoryCombo}s with their
+   * category options eagerly fetched in a single query. Used to prime the session before
+   * serialising category option combos, avoiding the per-combo N+1 select on the {@code
+   * categoryoptioncombos_categoryoptions} join table.
+   *
+   * @param categoryCombos the category combos whose option combos to load.
+   * @return the option combos with their category options initialised.
+   */
+  List<CategoryOptionCombo> getCategoryOptionCombosWithCategoryOptions(
+      Collection<CategoryCombo> categoryCombos);
+
   /**
    * Generates and persists a default Category, CategoryOption, CategoryCombo and
    * CategoryOptionCombo.

dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetService.java

@@ -29,9 +29,11 @@
  */
 package org.hisp.dhis.dataset;
 
+import java.util.Collection;
 import java.util.Date;
 import java.util.List;
 import org.apache.commons.collections4.SetValuedMap;
+import org.hisp.dhis.dataelement.DataElement;
 import org.hisp.dhis.dataentryform.DataEntryForm;
 import org.hisp.dhis.organisationunit.OrganisationUnit;
 import org.hisp.dhis.period.Period;
@@ -83,6 +85,19 @@ public interface DataSetService extends DataSetDataIntegrityProvider {
    */
   DataSet getDataSet(String uid);
 
+  /**
+   * Returns the distinct {@link DataElement}s that are members of the given data sets, loaded in a
+   * single query (with each data element's {@code dataSetElements} and {@code categoryCombo}
+   * eagerly fetched) to avoid N+1 selects when iterating {@link DataSet#getDataElements()} and
+   * {@link DataElement#getCategoryCombos()} across many data sets. Does not apply sharing/ACL
+   * predicates on the data elements; callers are responsible for authorizing access to the data
+   * sets.
+   *
+   * @param dataSets the data sets.
+   * @return the distinct data elements of the given data sets.
+   */
+  List<DataElement> getDataElementsByDataSet(Collection<DataSet> dataSets);
+
   /**
    * Returns the DataSet with the given UID. Bypasses the ACL system.
    *

dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetStore.java

@@ -53,6 +53,22 @@ public interface DataSetStore
 
   List<DataSetElement> getDataSetElementsByDataElement(Collection<DataElement> dataElements);
 
+  /**
+   * Returns the distinct {@link DataElement}s that are members of the given data sets (via {@link
+   * DataSetElement}), loaded in a single query. Each data element's {@link
+   * DataElement#getDataSetElements()} collection and its {@link DataElement#getCategoryCombo()} are
+   * eagerly fetched as well, so that {@link DataSet#getDataElements()} and {@link
+   * DataElement#getCategoryCombos()} can be evaluated without triggering per-element (or
+   * per-category-combo) N+1 selects.
+   *
+   * <p>This method does <b>not</b> apply sharing/ACL predicates on the data elements. Callers are
+   * responsible for authorizing access to the data sets.
+   *
+   * @param dataSets the data sets.
+   * @return the distinct data elements of the given data sets.
+   */
+  List<DataElement> getDataElementsByDataSet(Collection<DataSet> dataSets);
+
   /** Get all DatSets that have any of the CategoryCombos. */
   List<DataSet> getByCategoryCombo(List<CategoryCombo> sources);
 

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/category/DefaultCategoryService.java

@@ -514,6 +514,13 @@ public CategoryOptionCombo getCategoryOptionCombo(
     return categoryOptionComboStore.getCategoryOptionCombo(categoryCombo, categoryOptions);
   }
 
+  @Override
+  @Transactional(readOnly = true)
+  public List<CategoryOptionCombo> getCategoryOptionCombosWithCategoryOptions(
+      Collection<CategoryCombo> categoryCombos) {
+    return categoryOptionComboStore.getCategoryOptionCombosWithCategoryOptions(categoryCombos);
+  }
+
   @Override
   @Transactional(readOnly = true)
   public List<CategoryOptionCombo> getAllCategoryOptionCombos() {

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/category/hibernate/HibernateCategoryOptionComboStore.java

@@ -104,6 +104,26 @@ public CategoryOptionCombo getCategoryOptionCombo(
     return categoryOptionCombo;
   }
 
+  @Override
+  public List<CategoryOptionCombo> getCategoryOptionCombosWithCategoryOptions(
+      Collection<CategoryCombo> categoryCombos) {
+    if (categoryCombos == null || categoryCombos.isEmpty()) {
+      return List.of();
+    }
+
+    String hql =
+        """
+        select distinct coc from CategoryCombo cc
+        join cc.optionCombos coc
+        left join fetch coc.categoryOptions
+        where cc in :categoryCombos
+        """;
+
+    return getQuery(hql, CategoryOptionCombo.class)
+        .setParameter("categoryCombos", categoryCombos)
+        .list();
+  }
+
   @Override
   public void updateNames() {
     List<CategoryOptionCombo> categoryOptionCombos =

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java

@@ -41,6 +41,7 @@
 import lombok.RequiredArgsConstructor;
 import org.apache.commons.collections4.SetValuedMap;
 import org.hisp.dhis.association.jdbc.JdbcOrgUnitAssociationsStore;
+import org.hisp.dhis.dataelement.DataElement;
 import org.hisp.dhis.dataentryform.DataEntryForm;
 import org.hisp.dhis.organisationunit.OrganisationUnit;
 import org.hisp.dhis.period.Period;
@@ -99,6 +100,12 @@ public DataSet getDataSet(String uid) {
     return dataSetStore.getByUid(uid);
   }
 
+  @Override
+  @Transactional(readOnly = true)
+  public List<DataElement> getDataElementsByDataSet(Collection<DataSet> dataSets) {
+    return dataSetStore.getDataElementsByDataSet(dataSets);
+  }
+
   @Override
   @Transactional(readOnly = true)
   public DataSet getDataSetNoAcl(String uid) {

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/hibernate/HibernateDataSetStore.java

@@ -127,6 +127,25 @@ public List<DataSetElement> getDataSetElementsByDataElement(
         .list();
   }
 
+  @Override
+  public List<DataElement> getDataElementsByDataSet(Collection<DataSet> dataSets) {
+    if (dataSets == null || dataSets.isEmpty()) {
+      return List.of();
+    }
+
+    @Language("hql")
+    String hql =
+        """
+        select distinct de from DataSetElement dse
+        join dse.dataElement de
+        left join fetch de.dataSetElements
+        left join fetch de.categoryCombo
+        where dse.dataSet in :dataSets
+        """;
+
+    return getQuery(hql, DataElement.class).setParameter("dataSets", dataSets).list();
+  }
+
   @Override
   public List<DataSet> getByCategoryCombo(List<CategoryCombo> categoryCombos) {
     if (categoryCombos == null || categoryCombos.isEmpty()) return List.of();

dhis-2/dhis-services/dhis-service-dxf2/src/main/java/org/hisp/dhis/dxf2/metadata/DefaultDataSetMetadataExportService.java

@@ -144,12 +144,18 @@ public ObjectNode getDataSetMetadata() {
         dataSetService.getDataSetOrganisationUnitsAssociations();
 
     List<DataSet> dataSets = idObjectManager.getDataWriteAll(DataSet.class);
-    List<DataElement> dataElements = sortById(flatMapToSet(dataSets, DataSet::getDataElements));
+    List<DataElement> dataElements =
+        sortById(new HashSet<>(dataSetService.getDataElementsByDataSet(dataSets)));
     List<Indicator> indicators = sortById(flatMapToSet(dataSets, DataSet::getIndicators));
     List<CategoryCombo> dataElementCategoryCombos =
         sortById(flatMapToSet(dataElements, DataElement::getCategoryCombos));
     List<CategoryCombo> dataSetCategoryCombos =
         sortById(mapToSet(dataSets, DataSet::getCategoryCombo));
+
+    // Preload the category options of the data-element category-option-combos in a single query, so
+    // field-filter serialisation of categoryOptionCombos[...,categoryOptions~pluck[id]] does not
+    // trigger a per-combo N+1 (categoryoptioncombos_categoryoptions).
+    categoryService.getCategoryOptionCombosWithCategoryOptions(dataElementCategoryCombos);
     List<Category> dataElementCategories =
         sortById(flatMapToSet(dataElementCategoryCombos, CategoryCombo::getCategories));
     List<Category> dataSetCategories =

dhis-2/dhis-support/dhis-support-test/src/main/java/org/hisp/dhis/test/config/QueryCountDataSourceProxy.java

@@ -29,10 +29,17 @@
  */
 package org.hisp.dhis.test.config;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Locale;
 import javax.annotation.Nonnull;
 import javax.sql.DataSource;
+import net.ttddyy.dsproxy.ExecutionInfo;
+import net.ttddyy.dsproxy.QueryInfo;
 import net.ttddyy.dsproxy.listener.ChainListener;
 import net.ttddyy.dsproxy.listener.DataSourceQueryCountListener;
+import net.ttddyy.dsproxy.listener.QueryExecutionListener;
 import net.ttddyy.dsproxy.support.ProxyDataSourceBuilder;
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanPostProcessor;
@@ -57,10 +64,37 @@
  * <pre>{@code
  * assertDeleteCount(1);
  * }</pre>
+ *
+ * <p>In addition to the aggregate counts fed to {@code QueryCountHolder}, the executed SQL text is
+ * captured so a test can assert how many statements match a specific pattern (e.g. a join table)
+ * rather than just the total by type:
+ *
+ * <pre>{@code
+ * QueryCountDataSourceProxy.clearCapturedSql();
+ * // ... run the code under test ...
+ * assertTrue(QueryCountDataSourceProxy.countCapturedSqlMatching("some_join_table") <= 1);
+ * }</pre>
  */
 @Component
 public class QueryCountDataSourceProxy implements BeanPostProcessor {
 
+  private static final List<String> CAPTURED_SQL = Collections.synchronizedList(new ArrayList<>());
+
+  /** Clears the captured SQL. Call before the code under test. */
+  public static void clearCapturedSql() {
+    CAPTURED_SQL.clear();
+  }
+
+  /** Number of executed statements whose SQL contains the given (case-insensitive) pattern. */
+  public static long countCapturedSqlMatching(String pattern) {
+    String needle = pattern.toLowerCase(Locale.ROOT);
+    synchronized (CAPTURED_SQL) {
+      return CAPTURED_SQL.stream()
+          .filter(sql -> sql.toLowerCase(Locale.ROOT).contains(needle))
+          .count();
+    }
+  }
+
   @Override
   public Object postProcessBeforeInitialization(@Nonnull Object bean, @Nonnull String beanName) {
     return bean;
@@ -72,6 +106,19 @@ public Object postProcessAfterInitialization(@Nonnull Object bean, @Nonnull Stri
     if (bean instanceof DataSource originalDataSource && beanName.equals("actualDataSource")) {
       ChainListener listener = new ChainListener();
       listener.addListener(new DataSourceQueryCountListener());
+      listener.addListener(
+          new QueryExecutionListener() {
+            @Override
+            public void beforeQuery(ExecutionInfo execInfo, List<QueryInfo> queryInfoList) {
+              // No-op: SQL is recorded only after execution (afterQuery), so statements that error
+              // out or never run are not counted.
+            }
+
+            @Override
+            public void afterQuery(ExecutionInfo execInfo, List<QueryInfo> queryInfoList) {
+              queryInfoList.forEach(q -> CAPTURED_SQL.add(q.getQuery()));
+            }
+          });
       return ProxyDataSourceBuilder.create(originalDataSource)
           .name("query-count-datasource-proxy")
           .listener(listener)