chore: use enrollment service in TE aggregate

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/enrollment/DefaultEnrollmentService.java

@@ -243,6 +243,7 @@ public List<Enrollment> getEnrollments(@Nonnull Set<UID> uids) throws ForbiddenE
         .toList();
   }
 
+  @Nonnull
   @Override
   public List<Enrollment> getEnrollments(@Nonnull EnrollmentOperationParams params)
       throws ForbiddenException, BadRequestException {
@@ -255,6 +256,7 @@ public List<Enrollment> getEnrollments(@Nonnull EnrollmentOperationParams params
         queryParams.getOrganisationUnitMode());
   }
 
+  @Nonnull
   @Override
   public Page<Enrollment> getEnrollments(
       @Nonnull EnrollmentOperationParams params, PageParams pageParams)

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/HibernateTrackedEntityStore.java

@@ -182,13 +182,13 @@ public Page<TrackedEntityIdentifiers> getTrackedEntityIds(
   }
 
   private Page<TrackedEntityIdentifiers> getPage(
-      PageParams pageParams, List<TrackedEntityIdentifiers> teUids, LongSupplier enrollmentCount) {
+      PageParams pageParams, List<TrackedEntityIdentifiers> ids, LongSupplier enrollmentCount) {
     if (pageParams.isPageTotal()) {
       return Page.withTotals(
-          teUids, pageParams.getPage(), pageParams.getPageSize(), enrollmentCount.getAsLong());
+          ids, pageParams.getPage(), pageParams.getPageSize(), enrollmentCount.getAsLong());
     }
 
-    return Page.withoutTotals(teUids, pageParams.getPage(), pageParams.getPageSize());
+    return Page.withoutTotals(ids, pageParams.getPage(), pageParams.getPageSize());
   }
 
   @Override

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/TrackedEntityIdentifiers.java

@@ -28,7 +28,7 @@
 package org.hisp.dhis.tracker.export.trackedentity;
 
 /**
- * Temporary solution: tuple of primary key and uid needed by the aggregate store.
+ * Temporary solution: pair of primary key and uid needed by the aggregate store.
  *
  * @deprecated do not use this class! This is a temporary solution that will be removed.
  */

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/AsyncUtils.java

@@ -58,15 +58,4 @@ static <T> CompletableFuture<Multimap<String, T>> conditionalAsyncFetch(
         ? supplyAsync(supplier, executor)
         : supplyAsync(ArrayListMultimap::create, executor));
   }
-
-  /**
-   * Executes the Supplier asynchronously using the thread pool from the provided {@see Executor}
-   *
-   * @param supplier The Supplier to execute
-   * @return A CompletableFuture with the result of the Supplier
-   */
-  static <T> CompletableFuture<Multimap<String, T>> asyncFetch(
-      Supplier<Multimap<String, T>> supplier, Executor executor) {
-    return supplyAsync(supplier, executor);
-  }
 }

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/Context.java

@@ -32,6 +32,7 @@
 import lombok.Value;
 import org.hisp.dhis.tracker.export.trackedentity.TrackedEntityParams;
 import org.hisp.dhis.tracker.export.trackedentity.TrackedEntityQueryParams;
+import org.hisp.dhis.user.UserDetails;
 
 /**
  * @author Luciano Fiandesio
@@ -48,6 +49,8 @@ class Context {
   /** The current user uid */
   String userUid;
 
+  UserDetails userDetails;
+
   /** A list of group ID to which the user belongs */
   List<String> userGroups;
 

dhis-2/dhis-services/dhis-service-tracker/src/main/java/org/hisp/dhis/tracker/export/trackedentity/aggregates/EnrollmentAggregate.java

@@ -27,26 +27,20 @@
  */
 package org.hisp.dhis.tracker.export.trackedentity.aggregates;
 
-import static java.util.concurrent.CompletableFuture.allOf;
-import static org.hisp.dhis.tracker.export.trackedentity.aggregates.AsyncUtils.asyncFetch;
-import static org.hisp.dhis.tracker.export.trackedentity.aggregates.AsyncUtils.conditionalAsyncFetch;
-import static org.hisp.dhis.tracker.export.trackedentity.aggregates.ThreadPoolManager.getPool;
-
+import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.Multimap;
-import java.util.ArrayList;
-import java.util.HashSet;
-import java.util.LinkedHashSet;
 import java.util.List;
-import java.util.concurrent.CompletableFuture;
-import javax.annotation.Nonnull;
 import lombok.RequiredArgsConstructor;
-import org.hisp.dhis.note.Note;
+import org.hisp.dhis.common.UID;
+import org.hisp.dhis.feedback.BadRequestException;
+import org.hisp.dhis.feedback.ForbiddenException;
+import org.hisp.dhis.feedback.NotFoundException;
+import org.hisp.dhis.organisationunit.OrganisationUnit;
 import org.hisp.dhis.program.Enrollment;
-import org.hisp.dhis.program.Event;
-import org.hisp.dhis.relationship.RelationshipItem;
-import org.hisp.dhis.trackedentityattributevalue.TrackedEntityAttributeValue;
+import org.hisp.dhis.tracker.export.enrollment.EnrollmentOperationParams;
 import org.hisp.dhis.tracker.export.enrollment.EnrollmentService;
-import org.springframework.beans.factory.annotation.Qualifier;
+import org.hisp.dhis.tracker.export.trackedentity.TrackedEntityIdentifiers;
+import org.hisp.dhis.user.AuthenticationService;
 import org.springframework.stereotype.Component;
 
 /**
@@ -55,83 +49,54 @@
 @Component("org.hisp.dhis.tracker.trackedentity.aggregates.EnrollmentAggregate")
 @RequiredArgsConstructor
 class EnrollmentAggregate {
+  private final AuthenticationService authenticationService;
   private final EnrollmentService enrollmentService;
 
-  @Qualifier("org.hisp.dhis.tracker.trackedentity.aggregates.EnrollmentStore")
-  @Nonnull
-  private final EnrollmentStore enrollmentStore;
-
-  @Qualifier("org.hisp.dhis.tracker.trackedentity.aggregates.EventAggregate")
-  @Nonnull
-  private final EventAggregate eventAggregate;
-
   /**
    * Key: te uid , value Enrollment
    *
    * @param ids a List of {@see TrackedEntity} Primary Keys
    * @return a MultiMap where key is a {@see TrackedEntity} uid and the key a List of {@see
    *     Enrollment} objects
    */
-  Multimap<String, Enrollment> findByTrackedEntityIds(List<Long> ids, Context ctx) {
-    Multimap<String, Enrollment> enrollments =
-        enrollmentStore.getEnrollmentsByTrackedEntityIds(ids, ctx);
-
-    if (enrollments.isEmpty()) {
-      return enrollments;
+  Multimap<String, Enrollment> findByTrackedEntityIds(
+      List<TrackedEntityIdentifiers> ids, Context ctx) {
+    Multimap<String, Enrollment> result = ArrayListMultimap.create();
+
+    try {
+      authenticationService.obtainAuthentication(ctx.getUserUid());
+      ids.forEach(
+          id -> {
+            EnrollmentOperationParams params =
+                EnrollmentOperationParams.builder()
+                    .trackedEntity(UID.of(id.uid()))
+                    .trackedEntityType(ctx.getQueryParams().getTrackedEntityType())
+                    .enrollmentParams(ctx.getParams().getEnrollmentParams())
+                    .orgUnitMode(ctx.getQueryParams().getOrgUnitMode())
+                    .orgUnits(ctx.getQueryParams().getOrgUnits().toArray(new OrganisationUnit[0]))
+                    .program(ctx.getQueryParams().getProgram())
+                    .programStartDate(ctx.getQueryParams().getProgramEnrollmentStartDate())
+                    .programEndDate(ctx.getQueryParams().getProgramEnrollmentEndDate())
+                    .enrollmentStatus(ctx.getQueryParams().getEnrollmentStatus())
+                    .followUp(ctx.getQueryParams().getFollowUp())
+                    .build();
+            try {
+              result.putAll(id.uid(), enrollmentService.getEnrollments(params));
+            } catch (BadRequestException e) {
+              throw new IllegalArgumentException(
+                  "this must be a bug in how the EnrollmentOperationParams are built");
+            } catch (ForbiddenException e) {
+              // ForbiddenExceptions are caused when mapping the EnrollmentOperationParams. These
+              // params should already have been validated as they are coming from the
+              // TrackedEntityQueryParams. Other reasons the user does not have access to data will
+              // not be shown as such items are simply not returned in collections.
+            }
+          });
+    } catch (NotFoundException e) {
+      throw new RuntimeException(e);
+    } finally {
+      authenticationService.clearAuthentication();
     }
-
-    List<Long> enrollmentIds = enrollments.values().stream().map(Enrollment::getId).toList();
-
-    final CompletableFuture<Multimap<String, Event>> eventAsync =
-        conditionalAsyncFetch(
-            ctx.getParams().getEnrollmentParams().isIncludeEvents(),
-            () -> eventAggregate.findByEnrollmentIds(enrollmentIds, ctx),
-            getPool());
-
-    final CompletableFuture<Multimap<String, RelationshipItem>> relationshipAsync =
-        conditionalAsyncFetch(
-            ctx.getParams().getEnrollmentParams().isIncludeRelationships(),
-            () -> enrollmentStore.getRelationships(enrollmentIds, ctx),
-            getPool());
-
-    final CompletableFuture<Multimap<String, Note>> notesAsync =
-        asyncFetch(() -> enrollmentStore.getNotes(enrollmentIds), getPool());
-
-    final CompletableFuture<Multimap<String, TrackedEntityAttributeValue>> attributesAsync =
-        conditionalAsyncFetch(
-            ctx.getParams().getTeEnrollmentParams().isIncludeAttributes(),
-            () -> enrollmentStore.getAttributes(enrollmentIds, ctx),
-            getPool());
-
-    return allOf(eventAsync, notesAsync, relationshipAsync, attributesAsync)
-        .thenApplyAsync(
-            fn -> {
-              Multimap<String, Event> events = eventAsync.join();
-              Multimap<String, Note> notes = notesAsync.join();
-              Multimap<String, RelationshipItem> relationships = relationshipAsync.join();
-              Multimap<String, TrackedEntityAttributeValue> attributes = attributesAsync.join();
-
-              for (Enrollment enrollment : enrollments.values()) {
-                if (ctx.getParams().getTeEnrollmentParams().isIncludeEvents()) {
-                  enrollment.setEvents(new HashSet<>(events.get(enrollment.getUid())));
-                }
-                if (ctx.getParams().getTeEnrollmentParams().isIncludeRelationships()) {
-                  enrollment.setRelationshipItems(
-                      new HashSet<>(relationships.get(enrollment.getUid())));
-                }
-                if (ctx.getParams().getTeEnrollmentParams().isIncludeAttributes()) {
-                  enrollment
-                      .getTrackedEntity()
-                      .setTrackedEntityAttributeValues(
-                          new LinkedHashSet<>(attributes.get(enrollment.getUid())));
-                }
-
-                enrollment.setNotes(new ArrayList<>(notes.get(enrollment.getUid())));
-              }
-
-              return enrollments;
-            },
-            getPool())
-        .join();
+    return result;
   }
 }