feat: Batch entity when persisting [DHIS2-21378]

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

@@ -43,6 +43,7 @@
 import java.util.ArrayList;
 import java.util.Date;
 import java.util.EnumSet;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -79,8 +80,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 +159,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 +191,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 +264,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 void stageUpdate(V 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 //
@@ -407,19 +446,24 @@ protected void handleTrackedEntityAttributeValues(
     // already attached to the session, preventing a duplicate em.persist that would throw
     // EntityExistsException. The batch overlay catches the within-persister case (e.g. two
     // enrollments under the same TE both carrying the same attribute) where the second occurrence
-    // would otherwise produce a fresh instance with the same composite key.
+    // would otherwise produce a fresh instance with the same composite key. A transient TE
+    // (id == 0) has no DB rows yet -- the insert is still staged in the batch -- so binding it as
+    // a query parameter would throw TransientObjectException; skip the JPQL in that case and rely
+    // on the batch overlay alone.
     Map<MetadataIdentifier, TrackedEntityAttributeValue> attributeValueById =
-        entityManager
-            .createQuery(
-                "select v from TrackedEntityAttributeValue v where v.trackedEntity = :te",
-                TrackedEntityAttributeValue.class)
-            .setParameter("te", trackedEntity)
-            .getResultList()
-            .stream()
-            .collect(
-                Collectors.toMap(
-                    teav -> idSchemes.toMetadataIdentifier(teav.getAttribute()),
-                    Function.identity()));
+        trackedEntity.getId() == 0
+            ? new HashMap<>()
+            : entityManager
+                .createQuery(
+                    "select v from TrackedEntityAttributeValue v where v.trackedEntity = :te",
+                    TrackedEntityAttributeValue.class)
+                .setParameter("te", trackedEntity)
+                .getResultList()
+                .stream()
+                .collect(
+                    Collectors.toMap(
+                        teav -> idSchemes.toMetadataIdentifier(teav.getAttribute()),
+                        Function.identity()));
     batch
         .stagedFor(trackedEntity)
         .forEach(

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

@@ -33,26 +33,89 @@
 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: inserts and updates for TrackedEntity, Enrollment, TrackerEvent, SingleEvent;
+ * inserts only for Relationship (updates are rejected upstream by {@link
+ * AbstractTrackerPersister}); inserts, updates, and deletes for TEAVs. 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,25 +142,55 @@ Stream<TrackedEntityAttributeValue> stagedFor(TrackedEntity trackedEntity) {
   }
 
   Mark mark() {
-    return new Mark(teavInserts.size(), teavUpdates.size(), teavDeletes.size());
+    return new Mark(
+        new Mark.EntityCounts(teInserts.size(), teUpdates.size()),
+        new Mark.EntityCounts(enrollmentInserts.size(), enrollmentUpdates.size()),
+        new Mark.EntityCounts(trackerEventInserts.size(), trackerEventUpdates.size()),
+        new Mark.EntityCounts(singleEventInserts.size(), singleEventUpdates.size()),
+        relationshipInserts.size(),
+        new Mark.TeavCounts(teavInserts.size(), teavUpdates.size(), teavDeletes.size()));
   }
 
   void rollbackTo(Mark mark) {
-    truncate(teavInserts, mark.teavInsertSize);
-    truncate(teavUpdates, mark.teavUpdateSize);
-    truncate(teavDeletes, mark.teavDeleteSize);
+    truncate(teInserts, mark.te().inserts());
+    truncate(teUpdates, mark.te().updates());
+    truncate(enrollmentInserts, mark.enrollment().inserts());
+    truncate(enrollmentUpdates, mark.enrollment().updates());
+    truncate(trackerEventInserts, mark.trackerEvent().inserts());
+    truncate(trackerEventUpdates, mark.trackerEvent().updates());
+    truncate(singleEventInserts, mark.singleEvent().inserts());
+    truncate(singleEventUpdates, mark.singleEvent().updates());
+    truncate(relationshipInserts, mark.relationshipInserts());
+    truncate(teavInserts, mark.teav().inserts());
+    truncate(teavUpdates, mark.teav().updates());
+    truncate(teavDeletes, mark.teav().deletes());
   }
 
   /**
    * 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 +201,63 @@ 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(
+      EntityCounts te,
+      EntityCounts enrollment,
+      EntityCounts trackerEvent,
+      EntityCounts singleEvent,
+      int relationshipInserts,
+      TeavCounts teav) {
+
+    record EntityCounts(int inserts, int updates) {}
+
+    record TeavCounts(int inserts, int updates, int deletes) {}
+  }
 }