fix: Remove N+1 query when validating unique attributes [DHIS2-22157](2.41)

dhis-2/dhis-api/src/main/java/org/hisp/dhis/trackedentityattributevalue/TrackedEntityAttributeValueService.java

@@ -29,6 +29,8 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import javax.annotation.Nonnull;
 import org.hisp.dhis.trackedentity.TrackedEntity;
 import org.hisp.dhis.trackedentity.TrackedEntityAttribute;
 import org.hisp.dhis.user.User;
@@ -103,12 +105,29 @@ List<TrackedEntityAttributeValue> getTrackedEntityAttributeValues(
   int getCountOfAssignedTrackedEntityAttributeValues(TrackedEntityAttribute attribute);
 
   /**
-   * Retrieve a list of {@link TrackedEntityAttributeValue} that matches the values and the tea
-   * present in uniqueAttributes
+   * Finds the stored values of a unique attribute that match any of the given values, in any org
+   * unit. Values are compared case-insensitively. Use it for attributes unique across the whole
+   * system.
    *
-   * @param uniqueAttributes A map that links a list of values to a TEA
-   * @return TrackedEntityAttributeValue list
+   * @param attribute the unique attribute
+   * @param values the values to look for
+   * @return the matching values, empty if none
+   */
+  @Nonnull
+  List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      @Nonnull TrackedEntityAttribute attribute, @Nonnull Set<String> values);
+
+  /**
+   * Finds the stored values of an org unit scoped unique attribute that match one of the given (org
+   * unit, value) pairs exactly: a value only matches when stored for a tracked entity in the org
+   * unit it is paired with. Values are compared case-insensitively.
+   *
+   * @param attribute the unique attribute, unique within an org unit
+   * @param valuesByOrgUnitId the values to look for, keyed by the id of the org unit they are
+   *     unique in
+   * @return the matching values, empty if none
    */
-  List<TrackedEntityAttributeValue> getUniqueAttributeByValues(
-      Map<TrackedEntityAttribute, List<String>> uniqueAttributes);
+  @Nonnull
+  List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      @Nonnull TrackedEntityAttribute attribute, @Nonnull Map<Long, Set<String>> valuesByOrgUnitId);
 }

dhis-2/dhis-api/src/main/java/org/hisp/dhis/trackedentityattributevalue/TrackedEntityAttributeValueStore.java

@@ -27,10 +27,10 @@
  */
 package org.hisp.dhis.trackedentityattributevalue;
 
-import java.util.Collection;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import org.hisp.dhis.common.GenericStore;
-import org.hisp.dhis.program.Program;
 import org.hisp.dhis.trackedentity.TrackedEntity;
 import org.hisp.dhis.trackedentity.TrackedEntityAttribute;
 
@@ -82,33 +82,11 @@ public interface TrackedEntityAttributeValueStore
    */
   List<TrackedEntityAttributeValue> get(TrackedEntityAttribute attribute);
 
-  /**
-   * Gets a list of {@link TrackedEntityAttributeValue} that matches the parameters
-   *
-   * @param attribute {@link TrackedEntityAttribute} to get value for
-   * @param values List of literal values
-   * @return list of {@link TrackedEntityAttributeValue}
-   */
-  List<TrackedEntityAttributeValue> get(
-      TrackedEntityAttribute attribute, Collection<String> values);
+  List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Set<String> values);
 
-  /**
-   * Gets a list of {@link TrackedEntityAttributeValue} that matches the parameters
-   *
-   * @param attribute {@link TrackedEntityAttribute} to get value for
-   * @param value literal value to find within the specified {@link TrackedEntityAttribute}
-   * @return list of {@link TrackedEntityAttributeValue}
-   */
-  List<TrackedEntityAttributeValue> get(TrackedEntityAttribute attribute, String value);
-
-  /**
-   * Retrieve attribute values of an instance by a program.
-   *
-   * @param instance the TrackedEntity
-   * @param program the Program.
-   * @return TrackedEntityAttributeValue list
-   */
-  List<TrackedEntityAttributeValue> get(TrackedEntity instance, Program program);
+  List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Map<Long, Set<String>> valuesByOrgUnitId);
 
   /**
    * Return the number of assigned {@link TrackedEntityAttributeValue}s to the {@link

dhis-2/dhis-api/src/main/java/org/hisp/dhis/trackedentityattributevalue/UniqueAttributeValueMatch.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2004-2022, 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.trackedentityattributevalue;
+
+import javax.annotation.CheckForNull;
+import javax.annotation.Nonnull;
+import org.hisp.dhis.common.UID;
+
+/**
+ * A stored value of a unique tracked entity attribute that matches a value being imported.
+ *
+ * @param trackedEntity the tracked entity owning the value
+ * @param value the stored value, as persisted (not normalized to lower case)
+ * @param orgUnitId the id of the org unit of the tracked entity owning the value, only set for an
+ *     attribute unique within an org unit
+ */
+public record UniqueAttributeValueMatch(
+    @Nonnull UID trackedEntity, @Nonnull String value, @CheckForNull Long orgUnitId) {}

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/trackedentityattributevalue/DefaultTrackedEntityAttributeValueService.java

@@ -31,7 +31,7 @@
 
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
+import java.util.Set;
 import lombok.RequiredArgsConstructor;
 import org.apache.commons.lang3.StringUtils;
 import org.hisp.dhis.changelog.ChangeLogType;
@@ -217,10 +217,21 @@ private boolean addFileValue(TrackedEntityAttributeValue value) {
 
   @Override
   @Transactional(readOnly = true)
-  public List<TrackedEntityAttributeValue> getUniqueAttributeByValues(
-      Map<TrackedEntityAttribute, List<String>> uniqueAttributes) {
-    return uniqueAttributes.entrySet().stream()
-        .flatMap(entry -> this.attributeValueStore.get(entry.getKey(), entry.getValue()).stream())
-        .collect(Collectors.toList());
+  public List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Set<String> values) {
+    if (values.isEmpty()) {
+      return List.of();
+    }
+    return attributeValueStore.getUniqueAttributeValues(attribute, values);
+  }
+
+  @Override
+  @Transactional(readOnly = true)
+  public List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Map<Long, Set<String>> valuesByOrgUnitId) {
+    if (valuesByOrgUnitId.isEmpty()) {
+      return List.of();
+    }
+    return attributeValueStore.getUniqueAttributeValues(attribute, valuesByOrgUnitId);
   }
 }

dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/trackedentityattributevalue/hibernate/HibernateTrackedEntityAttributeValueStore.java

@@ -27,18 +27,23 @@
  */
 package org.hisp.dhis.trackedentityattributevalue.hibernate;
 
-import java.util.Collection;
+import com.google.common.collect.Lists;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
-import java.util.stream.Collectors;
+import java.util.Map;
+import java.util.Set;
 import javax.persistence.EntityManager;
 import org.apache.commons.lang3.StringUtils;
 import org.hibernate.query.Query;
+import org.hisp.dhis.common.UID;
 import org.hisp.dhis.hibernate.HibernateGenericStore;
-import org.hisp.dhis.program.Program;
 import org.hisp.dhis.trackedentity.TrackedEntity;
 import org.hisp.dhis.trackedentity.TrackedEntityAttribute;
 import org.hisp.dhis.trackedentityattributevalue.TrackedEntityAttributeValue;
 import org.hisp.dhis.trackedentityattributevalue.TrackedEntityAttributeValueStore;
+import org.hisp.dhis.trackedentityattributevalue.UniqueAttributeValueMatch;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.stereotype.Repository;
@@ -50,6 +55,25 @@
 public class HibernateTrackedEntityAttributeValueStore
     extends HibernateGenericStore<TrackedEntityAttributeValue>
     implements TrackedEntityAttributeValueStore {
+  /**
+   * Upper bound of values (or org unit/value pairs) bound in one query. A scoped query binds up to
+   * three parameters per pair, which keeps it below the Postgres limit of 32767 bind parameters.
+   */
+  private static final int MAX_VALUES_PER_QUERY = 10_000;
+
+  /**
+   * Selects only the columns the uniqueness validation needs. Selecting entities instead would load
+   * one tracked entity per matching row (the association from {@code TrackedEntityAttributeValue}
+   * is eager) and the users referenced by each org unit.
+   */
+  private static final String SELECT_UNIQUE_VALUES =
+      """
+      select te.uid, v.plainValue, te.organisationUnit.id
+      from TrackedEntityAttributeValue v
+      join v.trackedEntity te
+      where v.attribute = :attribute
+      """;
+
   public HibernateTrackedEntityAttributeValueStore(
       EntityManager entityManager, JdbcTemplate jdbcTemplate, ApplicationEventPublisher publisher) {
     super(entityManager, jdbcTemplate, publisher, TrackedEntityAttributeValue.class, false);
@@ -106,44 +130,103 @@ public List<TrackedEntityAttributeValue> get(TrackedEntityAttribute attribute) {
     return getList(typedQuery);
   }
 
+  /**
+   * Finds the stored values of {@code attribute} matching any of {@code values} (case-insensitive),
+   * in any org unit.
+   */
   @Override
-  public List<TrackedEntityAttributeValue> get(
-      TrackedEntityAttribute attribute, Collection<String> values) {
-    String query =
-        " from TrackedEntityAttributeValue v where v.attribute =:attribute and lower(v.plainValue) in :values";
-
-    Query<TrackedEntityAttributeValue> typedQuery =
-        getQuery(query)
-            .setParameter("attribute", attribute)
-            .setParameter(
-                "values", values.stream().map(StringUtils::lowerCase).collect(Collectors.toList()));
-
-    return getList(typedQuery);
+  public List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Set<String> values) {
+    List<UniqueAttributeValueMatch> matches = new ArrayList<>();
+    for (List<String> partition : Lists.partition(lowerCase(values), MAX_VALUES_PER_QUERY)) {
+      String hql = SELECT_UNIQUE_VALUES + " and lower(v.plainValue) in (:values)";
+      getSession()
+          .createQuery(hql, Object[].class)
+          .setParameter("attribute", attribute)
+          .setParameterList("values", partition)
+          .list()
+          .forEach(
+              row ->
+                  matches.add(
+                      new UniqueAttributeValueMatch(
+                          UID.of((String) row[0]), (String) row[1], (Long) row[2])));
+    }
+    return matches;
   }
 
+  /**
+   * Finds the stored values of {@code attribute} matching exactly one of the (org unit, value)
+   * pairs in {@code valuesByOrgUnitId} (value compared case-insensitively).
+   *
+   * <p>The value and org unit {@code in} filters let Postgres use the value or the org unit index,
+   * whichever is more selective. The key filter on {@code lower(value)|orgUnitId} then drops the
+   * cross product combinations that were not requested, e.g. a value only sent for org unit A that
+   * exists in org unit B. Postgres hashes the key list, so it does not slow the query down.
+   */
   @Override
-  public List<TrackedEntityAttributeValue> get(TrackedEntityAttribute attribute, String value) {
-    String query =
-        " from TrackedEntityAttributeValue v where v.attribute =:attribute and lower(v.plainValue) like :value";
-
-    Query<TrackedEntityAttributeValue> typedQuery =
-        getQuery(query)
-            .setParameter("attribute", attribute)
-            .setParameter("value", StringUtils.lowerCase(value));
-
-    return getList(typedQuery);
+  public List<UniqueAttributeValueMatch> getUniqueAttributeValues(
+      TrackedEntityAttribute attribute, Map<Long, Set<String>> valuesByOrgUnitId) {
+    List<UniqueAttributeValueMatch> matches = new ArrayList<>();
+    for (Map<Long, Set<String>> partition : partition(valuesByOrgUnitId)) {
+      Set<String> values = new HashSet<>();
+      List<String> keys = new ArrayList<>();
+      for (Map.Entry<Long, Set<String>> entry : partition.entrySet()) {
+        for (String value : lowerCase(entry.getValue())) {
+          values.add(value);
+          keys.add(value + "|" + entry.getKey());
+        }
+      }
+
+      String hql =
+          SELECT_UNIQUE_VALUES
+              + """
+               and lower(v.plainValue) in (:values)
+               and te.organisationUnit.id in (:orgUnitIds)
+               and concat(lower(v.plainValue), '|', str(te.organisationUnit.id)) in (:keys)
+              """;
+      getSession()
+          .createQuery(hql, Object[].class)
+          .setParameter("attribute", attribute)
+          .setParameterList("values", values)
+          .setParameterList("orgUnitIds", partition.keySet())
+          .setParameterList("keys", keys)
+          .list()
+          .forEach(
+              row ->
+                  matches.add(
+                      new UniqueAttributeValueMatch(
+                          UID.of((String) row[0]), (String) row[1], (Long) row[2])));
+    }
+    return matches;
   }
 
-  @Override
-  public List<TrackedEntityAttributeValue> get(TrackedEntity trackedEntity, Program program) {
-    String query =
-        " from TrackedEntityAttributeValue v where v.trackedEntity =:trackedEntity and v.attribute.program =:program";
-
-    Query<TrackedEntityAttributeValue> typedQuery = getQuery(query);
-    typedQuery.setParameter("trackedEntity", trackedEntity);
-    typedQuery.setParameter("program", program);
+  /**
+   * Splits the pairs into chunks of at most {@link #MAX_VALUES_PER_QUERY} pairs, so every query
+   * stays below the Postgres bind parameter limit (values + org units + keys per query).
+   */
+  private static List<Map<Long, Set<String>>> partition(Map<Long, Set<String>> valuesByOrgUnitId) {
+    List<Map<Long, Set<String>>> partitions = new ArrayList<>();
+    Map<Long, Set<String>> current = new HashMap<>();
+    int pairs = 0;
+    for (Map.Entry<Long, Set<String>> entry : valuesByOrgUnitId.entrySet()) {
+      for (String value : entry.getValue()) {
+        if (pairs == MAX_VALUES_PER_QUERY) {
+          partitions.add(current);
+          current = new HashMap<>();
+          pairs = 0;
+        }
+        current.computeIfAbsent(entry.getKey(), k -> new HashSet<>()).add(value);
+        pairs++;
+      }
+    }
+    if (!current.isEmpty()) {
+      partitions.add(current);
+    }
+    return partitions;
+  }
 
-    return getList(typedQuery);
+  private static List<String> lowerCase(Set<String> values) {
+    return values.stream().map(StringUtils::lowerCase).distinct().toList();
   }
 
   @Override