feat: Batch entity when persisting [DHIS2-21378]

Author: enricocolasante

dhis-2/dhis-tracker/src/main/java/org/hisp/dhis/tracker/imports/bundle/persister/AbstractTrackerPersister.java

@@ -79,8 +79,12 @@
 import org.hisp.dhis.tracker.imports.programrule.engine.Notification;
 import org.hisp.dhis.tracker.imports.report.Entity;
 import org.hisp.dhis.tracker.imports.report.TrackerTypeReport;
+import org.hisp.dhis.tracker.model.Enrollment;
+import org.hisp.dhis.tracker.model.Relationship;
+import org.hisp.dhis.tracker.model.SingleEvent;
 import org.hisp.dhis.tracker.model.TrackedEntity;
 import org.hisp.dhis.tracker.model.TrackedEntityAttributeValue;
+import org.hisp.dhis.tracker.model.TrackerEvent;
 import org.hisp.dhis.user.UserDetails;
 import org.springframework.jdbc.datasource.DataSourceUtils;
 
@@ -154,7 +158,7 @@ public PersistResult persist(TrackerBundle bundle) {
           // Save or update the entity
           //
           if (isNew(bundle, trackerDto)) {
-            entityManager.persist(convertedDto);
+            stageInsert(convertedDto, batch);
             updateDataValues(
                 bundle.getPreheat(),
                 trackerDto,
@@ -186,7 +190,7 @@ public PersistResult persist(TrackerBundle bundle) {
                   bundle.getUser(),
                   changeLogs,
                   batch);
-              entityManager.merge(convertedDto);
+              stageUpdate(convertedDto, batch);
               typeReport.getStats().incUpdated();
               typeReport.addEntity(objectReport);
               bundle.addUpdatedTrackedEntities(getUpdatedTrackedEntities(convertedDto));
@@ -259,6 +263,40 @@ public PersistResult persist(TrackerBundle bundle) {
     return new PersistResult(typeReport, notifications);
   }
 
+  private void stageInsert(V convertedDto, EntityWriteBatch batch) {
+    if (convertedDto instanceof TrackedEntity te) {
+      batch.stageInsert(te);
+    } else if (convertedDto instanceof Enrollment e) {
+      batch.stageInsert(e);
+    } else if (convertedDto instanceof TrackerEvent e) {
+      batch.stageInsert(e);
+    } else if (convertedDto instanceof SingleEvent e) {
+      batch.stageInsert(e);
+    } else if (convertedDto instanceof Relationship r) {
+      batch.stageInsert(r);
+    } else {
+      throw new IllegalArgumentException(
+          "Unsupported entity type: " + convertedDto.getClass().getName());
+    }
+  }
+
+  // Relationships are not updated -- the persister branch above ignores update payloads before
+  // this is called -- so no Relationship case is needed here.
+  private static void stageUpdate(Object convertedDto, EntityWriteBatch batch) {
+    if (convertedDto instanceof TrackedEntity te) {
+      batch.stageUpdate(te);
+    } else if (convertedDto instanceof Enrollment e) {
+      batch.stageUpdate(e);
+    } else if (convertedDto instanceof TrackerEvent e) {
+      batch.stageUpdate(e);
+    } else if (convertedDto instanceof SingleEvent e) {
+      batch.stageUpdate(e);
+    } else {
+      throw new IllegalArgumentException(
+          "Unsupported entity type for update: " + convertedDto.getClass().getName());
+    }
+  }
+
   // // // // // // // //
   // // // // // // // //
   // TEMPLATE METHODS //

dhis-2/dhis-tracker/src/main/java/org/hisp/dhis/tracker/imports/bundle/persister/EntityWriteBatch.java

@@ -33,26 +33,87 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.stream.Stream;
+import org.hisp.dhis.tracker.model.Enrollment;
+import org.hisp.dhis.tracker.model.Relationship;
+import org.hisp.dhis.tracker.model.SingleEvent;
 import org.hisp.dhis.tracker.model.TrackedEntity;
 import org.hisp.dhis.tracker.model.TrackedEntityAttributeValue;
+import org.hisp.dhis.tracker.model.TrackerEvent;
 
 /**
  * Accumulates entity-level writes staged during the persist loop and applies them at a single flush
  * point. Mirrors {@link ChangeLogAccumulator} and shares the same mark/rollback contract for
  * per-entity error isolation in non-atomic mode.
  *
- * <p>Phase 3 scope: TEAV inserts/updates/deletes only. The flush implementation currently delegates
- * back to {@link EntityManager} so that Hibernate continues to drive persistence while the staging
- * shape is in place. Phase 6 replaces {@link #flush(EntityManager)} with a connection-based
- * implementation that emits multi-row INSERT / unnest UPDATE / tuple DELETE statements. Top-level
- * entities (TrackedEntity, Enrollment, events, relationships) are added by Phase 4.
+ * <p>Current scope: TrackedEntity, Enrollment, TrackerEvent, SingleEvent, Relationship inserts and
+ * updates, plus TEAV inserts/updates/deletes. The flush implementation delegates back to {@link
+ * EntityManager} so that Hibernate continues to drive persistence while the staging shape is in
+ * place. Phase 6 replaces {@link #flush(EntityManager)} with a connection-based implementation that
+ * emits multi-row INSERT / unnest UPDATE / tuple DELETE statements.
+ *
+ * <p>Each {@link AbstractTrackerPersister#persist} call creates its own batch, so in practice only
+ * one top-level entity type list is populated per batch (the persister's own type) alongside any
+ * TEAVs staged by its attribute-handling code.
  */
 class EntityWriteBatch {
 
+  private final List<TrackedEntity> teInserts = new ArrayList<>();
+  private final List<TrackedEntity> teUpdates = new ArrayList<>();
+
+  private final List<Enrollment> enrollmentInserts = new ArrayList<>();
+  private final List<Enrollment> enrollmentUpdates = new ArrayList<>();
+
+  private final List<TrackerEvent> trackerEventInserts = new ArrayList<>();
+  private final List<TrackerEvent> trackerEventUpdates = new ArrayList<>();
+
+  private final List<SingleEvent> singleEventInserts = new ArrayList<>();
+  private final List<SingleEvent> singleEventUpdates = new ArrayList<>();
+
+  private final List<Relationship> relationshipInserts = new ArrayList<>();
+
   private final List<TrackedEntityAttributeValue> teavInserts = new ArrayList<>();
   private final List<TrackedEntityAttributeValue> teavUpdates = new ArrayList<>();
   private final List<TrackedEntityAttributeValue> teavDeletes = new ArrayList<>();
 
+  void stageInsert(TrackedEntity trackedEntity) {
+    teInserts.add(trackedEntity);
+  }
+
+  void stageUpdate(TrackedEntity trackedEntity) {
+    teUpdates.add(trackedEntity);
+  }
+
+  void stageInsert(Enrollment enrollment) {
+    enrollmentInserts.add(enrollment);
+  }
+
+  void stageUpdate(Enrollment enrollment) {
+    enrollmentUpdates.add(enrollment);
+  }
+
+  void stageInsert(TrackerEvent event) {
+    trackerEventInserts.add(event);
+  }
+
+  void stageUpdate(TrackerEvent event) {
+    trackerEventUpdates.add(event);
+  }
+
+  void stageInsert(SingleEvent event) {
+    singleEventInserts.add(event);
+  }
+
+  void stageUpdate(SingleEvent event) {
+    singleEventUpdates.add(event);
+  }
+
+  /**
+   * Relationships are never updated -- {@link AbstractTrackerPersister} ignores update payloads.
+   */
+  void stageInsert(Relationship relationship) {
+    relationshipInserts.add(relationship);
+  }
+
   void stageTeavInsert(TrackedEntityAttributeValue value) {
     teavInserts.add(value);
   }
@@ -79,10 +140,31 @@ Stream<TrackedEntityAttributeValue> stagedFor(TrackedEntity trackedEntity) {
   }
 
   Mark mark() {
-    return new Mark(teavInserts.size(), teavUpdates.size(), teavDeletes.size());
+    return new Mark(
+        teInserts.size(),
+        teUpdates.size(),
+        enrollmentInserts.size(),
+        enrollmentUpdates.size(),
+        trackerEventInserts.size(),
+        trackerEventUpdates.size(),
+        singleEventInserts.size(),
+        singleEventUpdates.size(),
+        relationshipInserts.size(),
+        teavInserts.size(),
+        teavUpdates.size(),
+        teavDeletes.size());
   }
 
   void rollbackTo(Mark mark) {
+    truncate(teInserts, mark.teInsertSize);
+    truncate(teUpdates, mark.teUpdateSize);
+    truncate(enrollmentInserts, mark.enrollmentInsertSize);
+    truncate(enrollmentUpdates, mark.enrollmentUpdateSize);
+    truncate(trackerEventInserts, mark.trackerEventInsertSize);
+    truncate(trackerEventUpdates, mark.trackerEventUpdateSize);
+    truncate(singleEventInserts, mark.singleEventInsertSize);
+    truncate(singleEventUpdates, mark.singleEventUpdateSize);
+    truncate(relationshipInserts, mark.relationshipInsertSize);
     truncate(teavInserts, mark.teavInsertSize);
     truncate(teavUpdates, mark.teavUpdateSize);
     truncate(teavDeletes, mark.teavDeleteSize);
@@ -92,12 +174,27 @@ void rollbackTo(Mark mark) {
    * Applies all staged writes through the provided {@link EntityManager}. Temporary delegating
    * implementation; Phase 6 swaps this for JDBC multi-row statements that take a {@link
    * java.sql.Connection}.
+   *
+   * <p>Top-level entities are flushed before TEAVs so that Hibernate assigns ids (and inserts the
+   * rows, when {@code em.flush()} runs) before any TEAV that references them. The top-level order
+   * (TE, Enrollment, TrackerEvent, SingleEvent, Relationship) matches the persister call order
+   * enforced by {@code DefaultTrackerBundleService.commit()} and is FK-safe.
    */
   void flush(EntityManager entityManager) {
-    if (teavInserts.isEmpty() && teavUpdates.isEmpty() && teavDeletes.isEmpty()) {
+    if (isEmpty()) {
       return;
     }
 
+    persistAll(entityManager, teInserts);
+    mergeAll(entityManager, teUpdates);
+    persistAll(entityManager, enrollmentInserts);
+    mergeAll(entityManager, enrollmentUpdates);
+    persistAll(entityManager, trackerEventInserts);
+    mergeAll(entityManager, trackerEventUpdates);
+    persistAll(entityManager, singleEventInserts);
+    mergeAll(entityManager, singleEventUpdates);
+    persistAll(entityManager, relationshipInserts);
+
     for (TrackedEntityAttributeValue v : teavInserts) {
       entityManager.persist(v);
     }
@@ -108,16 +205,64 @@ void flush(EntityManager entityManager) {
       entityManager.remove(entityManager.contains(v) ? v : entityManager.merge(v));
     }
 
+    teInserts.clear();
+    teUpdates.clear();
+    enrollmentInserts.clear();
+    enrollmentUpdates.clear();
+    trackerEventInserts.clear();
+    trackerEventUpdates.clear();
+    singleEventInserts.clear();
+    singleEventUpdates.clear();
+    relationshipInserts.clear();
     teavInserts.clear();
     teavUpdates.clear();
     teavDeletes.clear();
   }
 
+  private boolean isEmpty() {
+    return teInserts.isEmpty()
+        && teUpdates.isEmpty()
+        && enrollmentInserts.isEmpty()
+        && enrollmentUpdates.isEmpty()
+        && trackerEventInserts.isEmpty()
+        && trackerEventUpdates.isEmpty()
+        && singleEventInserts.isEmpty()
+        && singleEventUpdates.isEmpty()
+        && relationshipInserts.isEmpty()
+        && teavInserts.isEmpty()
+        && teavUpdates.isEmpty()
+        && teavDeletes.isEmpty();
+  }
+
+  private static <T> void persistAll(EntityManager entityManager, List<T> entities) {
+    for (T entity : entities) {
+      entityManager.persist(entity);
+    }
+  }
+
+  private static <T> void mergeAll(EntityManager entityManager, List<T> entities) {
+    for (T entity : entities) {
+      entityManager.merge(entity);
+    }
+  }
+
   private static <T> void truncate(List<T> list, int size) {
     if (list.size() > size) {
       list.subList(size, list.size()).clear();
     }
   }
 
-  record Mark(int teavInsertSize, int teavUpdateSize, int teavDeleteSize) {}
+  record Mark(
+      int teInsertSize,
+      int teUpdateSize,
+      int enrollmentInsertSize,
+      int enrollmentUpdateSize,
+      int trackerEventInsertSize,
+      int trackerEventUpdateSize,
+      int singleEventInsertSize,
+      int singleEventUpdateSize,
+      int relationshipInsertSize,
+      int teavInsertSize,
+      int teavUpdateSize,
+      int teavDeleteSize) {}
 }