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

david-mackessy · web-flow · commit db41604529e5 · 2026-07-21T10:58:51.000+02:00
* perf: DataEntry metadata with fewer n+1s [DHIS2-21757](2.41)

* format
diff --git a/dhis-2/dhis-api/src/main/java/org/hisp/dhis/category/CategoryOptionComboStore.java b/dhis-2/dhis-api/src/main/java/org/hisp/dhis/category/CategoryOptionComboStore.java
@@ -27,6 +27,7 @@
  */
 package org.hisp.dhis.category;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import org.hisp.dhis.common.IdentifiableObjectStore;
@@ -39,6 +40,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);
diff --git a/dhis-2/dhis-api/src/main/java/org/hisp/dhis/category/CategoryService.java b/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.
diff --git a/dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetService.java b/dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetService.java
@@ -27,6 +27,7 @@
  */
 package org.hisp.dhis.dataset;
 
+import java.util.Collection;
 import java.util.Date;
 import java.util.List;
 import org.apache.commons.collections4.SetValuedMap;
@@ -85,6 +86,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.
    *
diff --git a/dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetStore.java b/dhis-2/dhis-api/src/main/java/org/hisp/dhis/dataset/DataSetStore.java
@@ -27,8 +27,10 @@
  */
 package org.hisp.dhis.dataset;
 
+import java.util.Collection;
 import java.util.List;
 import org.hisp.dhis.common.IdentifiableObjectStore;
+import org.hisp.dhis.dataelement.DataElement;
 import org.hisp.dhis.dataentryform.DataEntryForm;
 import org.hisp.dhis.period.PeriodType;
 
@@ -58,4 +60,20 @@ public interface DataSetStore
    * @return a list of DataSets.
    */
   List<DataSet> getDataSetsByDataEntryForm(DataEntryForm dataEntryForm);
+
+  /**
+   * 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);
 }
diff --git a/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/category/hibernate/HibernateCategoryOptionComboStore.java b/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/category/hibernate/HibernateCategoryOptionComboStore.java
@@ -27,6 +27,7 @@
  */
 package org.hisp.dhis.category.hibernate;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Set;
 import javax.persistence.EntityManager;
@@ -98,6 +99,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 =
diff --git a/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataelement/DefaultCategoryService.java b/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataelement/DefaultCategoryService.java
@@ -496,6 +496,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() {
diff --git a/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java b/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/DefaultDataSetService.java
@@ -108,6 +108,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) {
diff --git a/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/hibernate/HibernateDataSetStore.java b/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/dataset/hibernate/HibernateDataSetStore.java
@@ -30,12 +30,14 @@
 import static com.google.common.base.Preconditions.checkNotNull;
 
 import com.google.common.collect.Lists;
+import java.util.Collection;
 import java.util.List;
 import javax.annotation.Nonnull;
 import javax.persistence.EntityManager;
 import javax.persistence.criteria.CriteriaBuilder;
 import org.hibernate.query.Query;
 import org.hisp.dhis.common.hibernate.HibernateIdentifiableObjectStore;
+import org.hisp.dhis.dataelement.DataElement;
 import org.hisp.dhis.dataentryform.DataEntryForm;
 import org.hisp.dhis.dataset.DataSet;
 import org.hisp.dhis.dataset.DataSetStore;
@@ -124,4 +126,23 @@ public List<DataSet> getDataSetsByDataEntryForm(DataEntryForm dataEntryForm) {
   public List<DataSet> getDataSetsNotAssignedToOrganisationUnits() {
     return getQuery("from DataSet ds where size(ds.sources) = 0").list();
   }
+
+  @Override
+  public List<DataElement> getDataElementsByDataSet(Collection<DataSet> dataSets) {
+    if (dataSets == null || dataSets.isEmpty()) {
+      return List.of();
+    }
+
+    return getQuery(
+            """
+            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
+            """,
+            DataElement.class)
+        .setParameter("dataSets", dataSets)
+        .list();
+  }
 }
diff --git a/dhis-2/dhis-services/dhis-service-dxf2/src/main/java/org/hisp/dhis/dxf2/metadata/DefaultDataSetMetadataExportService.java b/dhis-2/dhis-services/dhis-service-dxf2/src/main/java/org/hisp/dhis/dxf2/metadata/DefaultDataSetMetadataExportService.java
@@ -145,12 +145,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 =
diff --git a/dhis-2/dhis-support/dhis-support-test/pom.xml b/dhis-2/dhis-support/dhis-support-test/pom.xml
@@ -99,6 +99,10 @@
       <groupId>org.hibernate</groupId>
       <artifactId>hibernate-core</artifactId>
     </dependency>
+    <dependency>
+      <groupId>net.ttddyy</groupId>
+      <artifactId>datasource-proxy</artifactId>
+    </dependency>
     <dependency>
       <groupId>javax.persistence</groupId>
       <artifactId>javax.persistence-api</artifactId>
diff --git a/dhis-2/dhis-support/dhis-support-test/src/main/java/org/hisp/dhis/test/config/QueryCountDataSourceProxy.java b/dhis-2/dhis-support/dhis-support-test/src/main/java/org/hisp/dhis/test/config/QueryCountDataSourceProxy.java
@@ -0,0 +1,127 @@
+/*
+ * Copyright (c) 2004-2024, University of Oslo
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ * Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ * Neither the name of the HISP project nor the names of its contributors may
+ * be used to endorse or promote products derived from this software without
+ * specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ * ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+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;
+import org.springframework.stereotype.Component;
+
+/**
+ * Test class that enhances the originalDatasource with the ability to count queries, and returns it
+ * as a proxy. This should only be used in tests. To use, simply add:
+ *
+ * <pre>{@code
+ * @ContextConfiguration(classes = {QueryCountDataSourceProxy.class})
+ * }</pre>
+ *
+ * on your test class and then use:
+ *
+ * <pre>{@code
+ * SQLStatementCountValidator.reset();
+ * }</pre>
+ *
+ * before the condition to test and then use asserts afterward e.g.
+ *
+ * <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;
+  }
+
+  @Override
+  public Object postProcessAfterInitialization(@Nonnull Object bean, @Nonnull String beanName)
+      throws BeansException {
+    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)
+          .build();
+    }
+    return bean;
+  }
+}
diff --git a/dhis-2/dhis-test-integration/pom.xml b/dhis-2/dhis-test-integration/pom.xml
@@ -184,6 +184,11 @@
       <artifactId>dhis-support-test</artifactId>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>net.ttddyy</groupId>
+      <artifactId>datasource-proxy</artifactId>
+      <scope>test</scope>
+    </dependency>
     <dependency>
       <groupId>org.hisp.dhis.parser</groupId>
       <artifactId>dhis-antlr-expression-parser</artifactId>
diff --git a/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/dxf2/metadata/DataSetMetadataExportServiceQueryCountTest.java b/dhis-2/dhis-test-integration/src/test/java/org/hisp/dhis/dxf2/metadata/DataSetMetadataExportServiceQueryCountTest.java
