From 768c8e218921b4267d1031bd810a1c39341dfde7 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Fri, 18 Sep 2026 11:41:51 -0400 Subject: [PATCH] Implement more efficient bulk put/update in JPATM This is more complicated than the bulk load unfortunately. Previously, for update/put or their plural equivalents, we'd have to run one or two SELECT statements per entity to check if it exists already (and for Hibernate to load the object into the context to merge it with the input). Both updateAll and putAll now run, per input class, one large SELECT statement to find out which entities already exist, which also loads the entities into the persistence context (which is basically a cache). For updateAll, we require that all of the entities already exist. For putAll, we have a split path - for entities that already exist, we just call a normal merge() - for entities that don't already exist, we run merge() while instructing Hibernate that the object is transient (i.e. new). This means that we get the "put" behavior (not modifying the original object) while avoid a superfluous SELECT statement that'd check again to see if the object exists in the DB already. For new objects in putAll, we also use EntityManager::merge (rather than insert/persist) so that Hibernate creates a deep copy of the entity instead of mutating the caller's instance in place. To prevent Hibernate's DefaultMergeEventListener from firing a redundant SELECT statement when merging a known-new entity with a non-null ID, TransactionInfo implements Interceptor and overrides isTransient(Object) to return Boolean.TRUE while merging a known-transient entity. We add comments clarifying that "insert" will modify an entity in place (i.e. UpdateAutoTimestamp or generated IDs) whereas "put" will not. This has always been the case; the comments are just to clarify. --- .../JpaTransactionManagerImpl.java | 190 +++++++++++++++--- .../transaction/TransactionManager.java | 64 +++++- .../ResaveAllEppResourcesPipelineTest.java | 19 +- .../model/common/DnsRefreshRequestTest.java | 11 + .../JpaTransactionManagerImplTest.java | 117 ++++++++++- 5 files changed, 354 insertions(+), 47 deletions(-) diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java index 2895380e008..d2c0704c14f 100644 --- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java +++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java @@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Multimaps; +import com.google.common.collect.Sets; import com.google.common.collect.Streams; import com.google.common.flogger.FluentLogger; import com.google.common.flogger.StackSize; @@ -50,6 +51,7 @@ import jakarta.persistence.LockModeType; import jakarta.persistence.Parameter; import jakarta.persistence.PersistenceException; +import jakarta.persistence.PersistenceUnitUtil; import jakarta.persistence.Query; import jakarta.persistence.TemporalType; import jakarta.persistence.TypedQuery; @@ -79,7 +81,9 @@ import java.util.regex.Pattern; import java.util.stream.Stream; import javax.annotation.Nullable; +import org.hibernate.Interceptor; import org.hibernate.Session; +import org.hibernate.SessionBuilder; import org.hibernate.SessionFactory; import org.hibernate.cfg.Environment; @@ -253,20 +257,17 @@ public T transactNoRetry( } TransactionInfo txnInfo = transactionInfo.get(); - txnInfo.entityManager = - logSqlStatements - ? emf.unwrap(SessionFactory.class) - .withOptions() - .statementInspector( - new UnaryOperator() { - @Override - public String apply(String s) { - logger.atInfo().log(SQL_STATEMENT_LOG_SENTINEL_FORMAT, s); - return s; - } - }) - .openSession() - : emf.createEntityManager(); + SessionBuilder sessionBuilder = + emf.unwrap(SessionFactory.class).withOptions().interceptor(txnInfo); + if (logSqlStatements) { + sessionBuilder.statementInspector( + (UnaryOperator) + s -> { + logger.atInfo().log(SQL_STATEMENT_LOG_SENTINEL_FORMAT, s); + return s; + }); + } + txnInfo.entityManager = sessionBuilder.openSession(); if (readOnly) { // Disable Hibernate's dirty object check on flushing, it has become more aggressive in v6. txnInfo.entityManager.unwrap(Session.class).setDefaultReadOnly(true); @@ -359,16 +360,20 @@ public Instant getTxTime() { } /** - * Inserts an object into the database. + * Inserts a new object into the database, throwing an exception if it already exists. + * + *

This method delegates to {@link EntityManager#persist} and therefore modifies {@code + * entity} in place (e.g., assigning {@link jakarta.persistence.GeneratedValue} IDs, setting + * auto-timestamps, and wrapping collections). * *

If {@code entity} has an auto-generated identity field (i.e., a field annotated with {@link * jakarta.persistence.GeneratedValue}), the caller must not assign a value to this field, * otherwise Hibernate would mistake the entity as detached and raise an error. * *

The practical implication of the above is that when inserting such an entity using a - * retriable transaction , the entity should be instantiated inside the transaction body. A failed - * attempt may still assign and ID to the entity, therefore reusing the same entity would cause - * retries to fail. + * retriable transaction, the entity should be instantiated inside the transaction body. A failed + * attempt may still assign an ID to the entity, therefore reusing the same entity would cause + * retries to fail; otherwise, prefer {@link #put}. */ @Override public void insert(Object entity) { @@ -386,9 +391,15 @@ public void insertAll(ImmutableCollection entities) { @Override public void insertAll(ImmutableObject... entities) { - insertAll(ImmutableSet.copyOf(entities)); + insertAll(ImmutableList.copyOf(entities)); } + /** + * Persists a new object or updates an existing object in the database. + * + *

Unlike {@link #insert}, this method delegates to {@link EntityManager#merge} to make a deep + * copy and never modifies {@code entity} in place. + */ @Override public void put(Object entity) { checkArgumentNotNull(entity, "entity must be specified"); @@ -399,32 +410,73 @@ public void put(Object entity) { @Override public void putAll(ImmutableObject... entities) { checkArgumentNotNull(entities, "entities must be specified"); - assertInTransaction(); - for (Object entity : entities) { - put(entity); - } + putAll(ImmutableList.copyOf(entities)); } @Override public void putAll(ImmutableCollection entities) { checkArgumentNotNull(entities, "entities must be specified"); assertInTransaction(); - entities.forEach(this::put); + // Group entities by class so we can batch-load existing entities per concrete type via + // Session::findMultiple. + ImmutableListMultimap, ?> entitiesByClass = + Multimaps.index(entities, Object::getClass); + // Pre-warm the Hibernate first-level cache (persistence context) in one batched SELECT per + // entity class and record which entities already exist in the database. Without this, calling + // EntityManager::merge on detached entities would execute an individual SELECT per entity. + Set existingEntities = Collections.newSetFromMap(new IdentityHashMap<>()); + for (Class clazz : entitiesByClass.keySet()) { + existingEntities.addAll( + findExistingEntitiesAndLoadContext(clazz, entitiesByClass.get(clazz))); + } + for (Object entity : entities) { + if (existingEntities.contains(entity)) { + // The entity already exists in the persistence context (loaded from the DB above); + // updateObject merges into the pre-warmed persistence context without issuing a SELECT. + transactionInfo.get().updateObject(entity); + } else { + // The entity does not exist in the DB yet (either id == null or findMultiple returned + // null). By calling "merge" while explicitly telling Hibernate that the object is + // transient, we get a deep copy (so, not modifying the immutable original) while still + // avoiding issuing an additional SELECT. + transactionInfo.get().mergeTransientObject(entity); + } + } } @Override public void update(Object entity) { checkArgumentNotNull(entity, "entity must be specified"); - assertInTransaction(); - checkArgument(exists(entity), "Given entity does not exist"); - transactionInfo.get().updateObject(entity); + updateAll(ImmutableList.of(entity)); } @Override public void updateAll(ImmutableCollection entities) { checkArgumentNotNull(entities, "entities must be specified"); assertInTransaction(); - entities.forEach(this::update); + ImmutableListMultimap, ?> entitiesByClass = + Multimaps.index(entities, Object::getClass); + // Instead of merging the entities one by one, we load the entities in one go then merge them. + // If we go one-by-one, we execute one or two SELECT statements per entity (one to check + // existence, one to load into the persistence context). Instead, we run one large SELECT + // statement to check all the entities' existences. This has the extra benefit of loading all + // the entities into the persistence context (cache) at once instead of having to load the + // entities into the cache one by one. + for (Class clazz : entitiesByClass.keySet()) { + ImmutableList entitiesForThisClass = entitiesByClass.get(clazz); + ImmutableList existingEntities = + findExistingEntitiesAndLoadContext(clazz, entitiesForThisClass); + // All entities should be present and accounted for after loading the persistence context + if (existingEntities.size() != entitiesForThisClass.size()) { + throw new IllegalArgumentException( + String.format( + "Entity/entities passed to updateAll do not already exist: %s", + Sets.difference( + ImmutableSet.copyOf(entitiesForThisClass), + ImmutableSet.copyOf(existingEntities)))); + } + } + entities.forEach(transactionInfo.get()::updateObject); } @Override @@ -806,11 +858,54 @@ private T detach(@Nullable T entity) { return entity; } - private static class TransactionInfo { + /** + * Finds existing entities and pre-warms the persistence context (cache). + * + *

This first filters out entities that don't have IDs (i.e. new objects with generated IDs). + * Then we call Session::findMultiple to find existing entities while also populating the + * persistence context (cache). This means that later calls to merge() won't need to SELECT, as + * the object is already loaded. + */ + private ImmutableList findExistingEntitiesAndLoadContext( + Class clazz, ImmutableList entitiesWithThisClass) { + Session session = getEntityManager().unwrap(Session.class); + PersistenceUnitUtil persistenceUnitUtil = emf.getPersistenceUnitUtil(); + ImmutableList.Builder entitiesWithIdsBuilder = new ImmutableList.Builder<>(); + ImmutableList.Builder idsBuilder = new ImmutableList.Builder<>(); + for (Object entity : entitiesWithThisClass) { + Object id = persistenceUnitUtil.getIdentifier(entity); + // Skip entities with null IDs (e.g., unpersisted @GeneratedValue entities) since they + // cannot exist in the database yet. + if (id != null) { + entitiesWithIdsBuilder.add(entity); + idsBuilder.add(id); + } + } + ImmutableList entitiesWithIds = entitiesWithIdsBuilder.build(); + ImmutableList ids = idsBuilder.build(); + checkArgument( + ids.size() == ImmutableSet.copyOf(ids).size(), + "Multiple entities of type %s with the same ID", + clazz.getSimpleName()); + ImmutableList.Builder existingEntitiesBuilder = new ImmutableList.Builder<>(); + if (!ids.isEmpty()) { + // Session::findMultiple's return value includes nulls for missing entities + List loaded = session.findMultiple(clazz, ids); + for (int i = 0; i < ids.size(); i++) { + if (loaded.get(i) != null) { + existingEntitiesBuilder.add(entitiesWithIds.get(i)); + } + } + } + return existingEntitiesBuilder.build(); + } + + private static class TransactionInfo implements Interceptor { EntityManager entityManager; boolean inTransaction = false; Instant transactionTime; Supplier idProvider; + @Nullable Object transientEntityToMerge; // The set of entity objects that have been either persisted (via insert()) or merged (via // put()/update()). If the entity manager returns these as a result of a find() or query @@ -818,6 +913,12 @@ private static class TransactionInfo { // them to not be saved to the database -- so we throw an exception instead. Set objectsToSave = Collections.newSetFromMap(new IdentityHashMap<>()); + @Override + public Boolean isTransient(Object entity) { + // "null" indicates "we don't know if this is transient" + return (entity != null && entity == transientEntityToMerge) ? Boolean.TRUE : null; + } + /** Start a new transaction. */ private void start(Clock clock, Supplier idProvider) { checkArgumentNotNull(clock); @@ -830,6 +931,7 @@ private void clear() { idProvider = null; inTransaction = false; transactionTime = null; + transientEntityToMerge = null; objectsToSave = Collections.newSetFromMap(new IdentityHashMap<>()); if (entityManager != null) { // Close this EntityManager just let the connection pool be able to reuse it, it doesn't @@ -839,13 +941,39 @@ private void clear() { } } - /** Does the full "update" on an object including all internal housekeeping. */ + /** + * Merges {@code object} into the session (making a deep copy without modifying {@code object} + * in place) and records the managed copy for internal housekeeping. + */ private void updateObject(Object object) { Object merged = entityManager.merge(object); objectsToSave.add(merged); } - /** Does the full "insert" on a new object including all internal housekeeping. */ + /** + * Merges a known-transient (new) object into the session without firing a database SELECT. + * + *

Hibernate performs deep copies on objects when calling "merge" which we want in order to + * avoid modifying the original object. However, "merge" issues an additional SELECT statement + * unless we explicitly tell Hibernate the object is transient. By setting {@link + * #transientEntityToMerge}, {@link #isTransient(Object)} returns {@code Boolean.TRUE} when + * Hibernate's {@code DefaultMergeEventListener} inspects the entity, causing Hibernate to take + * the {@code entityIsTransient} path (deep-copying the entity and scheduling an INSERT) without + * executing a {@code SELECT} query even when the entity has a pre-assigned non-null ID. + */ + private void mergeTransientObject(Object object) { + transientEntityToMerge = object; + try { + updateObject(object); + } finally { + transientEntityToMerge = null; + } + } + + /** + * Persists a new {@code object} directly in the session (modifying {@code object} in place) and + * records it for internal housekeeping. + */ private void insertObject(Object object) { entityManager.persist(object); objectsToSave.add(object); diff --git a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java index 3bfdb61cf38..746a29aa720 100644 --- a/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java +++ b/core/src/main/java/google/registry/persistence/transaction/TransactionManager.java @@ -131,31 +131,77 @@ public interface TransactionManager { /** Returns the Instant associated with the start of this particular transaction attempt. */ Instant getTxTime(); - /** Persists a new entity in the database, throws exception if the entity already exists. */ + /** + * Persists a new entity in the database, throwing an exception if the entity already exists. + * + *

Note: This method uses {@link jakarta.persistence.EntityManager#persist} and modifies the + * input {@code entity} in place (e.g., assigning {@link jakarta.persistence.GeneratedValue} + * IDs, setting auto-timestamps, and wrapping collections). When used inside a retriable + * transaction, {@code entity} should be instantiated inside the transaction body so retries do + * not reuse a mutated instance; otherwise, prefer {@link #put}. + */ void insert(Object entity); - /** Persists all new entities in the database, throws exception if any entity already exists. */ + /** + * Persists all new entities in the database, throwing an exception if any entity already exists. + * + *

Like {@link #insert}, this method modifies the input entities in place. + */ void insertAll(ImmutableCollection entities); - /** Persists all new entities in the database, throws exception if any entity already exists. */ + /** + * Persists all new entities in the database, throwing an exception if any entity already exists. + * + *

Like {@link #insert}, this method modifies the input entities in place. + */ void insertAll(ImmutableObject... entities); - /** Persists a new entity or update the existing entity in the database. */ + /** + * Persists a new entity or updates an existing entity in the database. + * + *

Unlike {@link #insert}, this method uses {@link jakarta.persistence.EntityManager#merge} to + * make a deep copy and never modifies the input {@code entity} in place. + */ void put(Object entity); - /** Persists all new entities or updates the existing entities in the database. */ + /** + * Persists all new entities or updates existing entities in the database. + * + *

Unlike {@link #insertAll}, this method uses {@link jakarta.persistence.EntityManager#merge} + * to make a deep copy and never modifies the input entities in place. + */ void putAll(ImmutableObject... entities); - /** Persists all new entities or updates the existing entities in the database. */ + /** + * Persists all new entities or updates existing entities in the database. + * + *

Unlike {@link #insertAll}, this method uses {@link jakarta.persistence.EntityManager#merge} + * to make a deep copy and never modifies the input entities in place. + */ void putAll(ImmutableCollection entities); - /** Updates an entity in the database, throws exception if the entity does not exist. */ + /** + * Updates an existing entity in the database, throwing an exception if it does not exist. + * + *

This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and + * never modifies the input {@code entity} in place. + */ void update(Object entity); - /** Updates all entities in the database, throws exception if any entity does not exist. */ + /** + * Updates all existing entities in the database, throwing an exception if any does not exist. + * + *

This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and + * never modifies the input entities in place. + */ void updateAll(ImmutableCollection entities); - /** Updates all entities in the database, throws exception if any entity does not exist. */ + /** + * Updates all existing entities in the database, throwing an exception if any does not exist. + * + *

This method uses {@link jakarta.persistence.EntityManager#merge} to make a deep copy and + * never modifies the input entities in place. + */ void updateAll(ImmutableObject... entities); /** Returns whether the given entity with same ID exists. */ diff --git a/core/src/test/java/google/registry/beam/resave/ResaveAllEppResourcesPipelineTest.java b/core/src/test/java/google/registry/beam/resave/ResaveAllEppResourcesPipelineTest.java index ae4316f023f..c9d226ef09e 100644 --- a/core/src/test/java/google/registry/beam/resave/ResaveAllEppResourcesPipelineTest.java +++ b/core/src/test/java/google/registry/beam/resave/ResaveAllEppResourcesPipelineTest.java @@ -15,6 +15,7 @@ package google.registry.beam.resave; import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import static google.registry.testing.DatabaseHelper.createTld; @@ -27,10 +28,11 @@ import static google.registry.testing.DatabaseHelper.persistNewRegistrars; import static google.registry.util.DateTimeUtils.minusDays; import static google.registry.util.DateTimeUtils.plusYears; +import static org.mockito.Mockito.atLeastOnce; import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import com.google.common.collect.ImmutableCollection; import google.registry.beam.TestPipelineExtension; import google.registry.model.EppResource; import google.registry.model.domain.Domain; @@ -138,6 +140,7 @@ void testPipeline_expiredGracePeriod() { } @Test + @SuppressWarnings("unchecked") void testPipeline_fastOnlySavesChanged() { Instant now = fakeClock.now(); persistDomainWithDependentResources("renewed", "tld", now, now, plusYears(now, 1)); @@ -145,14 +148,16 @@ void testPipeline_fastOnlySavesChanged() { // Spy the transaction manager so we can be sure we're only saving the renewed domain JpaTransactionManager spy = spy(tm()); TransactionManagerFactory.setJpaTm(() -> spy); - ArgumentCaptor domainPutCaptor = ArgumentCaptor.forClass(Domain.class); + ArgumentCaptor> domainPutCaptor = + ArgumentCaptor.forClass(ImmutableCollection.class); runPipeline(); // We should only be attempting to put the one changed domain into the DB - verify(spy).put(domainPutCaptor.capture()); - assertThat(domainPutCaptor.getValue().getDomainName()).isEqualTo("renewed.tld"); + verify(spy).putAll(domainPutCaptor.capture()); + assertThat(getOnlyElement(domainPutCaptor.getValue()).getDomainName()).isEqualTo("renewed.tld"); } @Test + @SuppressWarnings("unchecked") void testPipeline_notFastResavesAll() { options.setFast(false); Instant now = fakeClock.now(); @@ -163,12 +168,14 @@ void testPipeline_notFastResavesAll() { // Spy the transaction manager so we can be sure we're attempting to save everything JpaTransactionManager spy = spy(tm()); TransactionManagerFactory.setJpaTm(() -> spy); - ArgumentCaptor eppResourcePutCaptor = ArgumentCaptor.forClass(EppResource.class); + ArgumentCaptor> eppResourcePutCaptor = + ArgumentCaptor.forClass(ImmutableCollection.class); runPipeline(); // We should be attempting to put both domains in, even the unchanged one - verify(spy, times(2)).put(eppResourcePutCaptor.capture()); + verify(spy, atLeastOnce()).putAll(eppResourcePutCaptor.capture()); assertThat( eppResourcePutCaptor.getAllValues().stream() + .flatMap(ImmutableCollection::stream) .map(EppResource::getRepoId) .collect(toImmutableSet())) .containsExactly(renewed.getRepoId(), nonRenewed.getRepoId()); diff --git a/core/src/test/java/google/registry/model/common/DnsRefreshRequestTest.java b/core/src/test/java/google/registry/model/common/DnsRefreshRequestTest.java index 41b0b44d1b4..ab82a988c77 100644 --- a/core/src/test/java/google/registry/model/common/DnsRefreshRequestTest.java +++ b/core/src/test/java/google/registry/model/common/DnsRefreshRequestTest.java @@ -49,6 +49,17 @@ void testPersistence() { assertThat(requests.get(0).id).isNotNull(); } + @Test + void testPutAll_leavesTransientIdNull() { + assertThat(request.id).isNull(); + tm().transact(() -> tm().putAll(request)); + assertThat(request.id).isNull(); + ImmutableList requests = loadAllOf(DnsRefreshRequest.class); + assertThat(requests).hasSize(1); + assertThat(requests.get(0).id).isNotNull(); + assertAboutImmutableObjects().that(requests.get(0)).isEqualExceptFields(request, "id"); + } + @Test void testNullValues() { // type diff --git a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java index 580fe28822e..18f5277fbef 100644 --- a/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java +++ b/core/src/test/java/google/registry/persistence/transaction/JpaTransactionManagerImplTest.java @@ -54,6 +54,8 @@ import jakarta.persistence.RollbackException; import java.io.Serializable; import java.util.NoSuchElementException; +import org.hibernate.SessionFactory; +import org.hibernate.stat.Statistics; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.api.function.Executable; @@ -480,6 +482,76 @@ void putAll_succeeds() { .containsExactlyElementsIn(moreEntities); } + @Test + void putAll_mixedEntityTypes_succeeds() { + persistResource(theEntity); + TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated"); + TestEntity newEntity = new TestEntity("newEntity", "new_data"); + tm().transact( + () -> tm().putAll(ImmutableList.of(updatedTheEntity, compoundIdEntity, newEntity))); + assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity); + assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey))) + .isEqualTo(compoundIdEntity); + assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity")))) + .isEqualTo(newEntity); + } + + @Test + void putAll_varargs_succeeds() { + persistResource(theEntity); + TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated"); + tm().transact(() -> tm().putAll(updatedTheEntity, compoundIdEntity)); + assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity); + assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey))) + .isEqualTo(compoundIdEntity); + } + + @Test + void putAll_duplicateIds_throws() { + assertThat( + assertThrows( + IllegalArgumentException.class, + () -> + tm().transact( + () -> + tm().putAll( + new TestEntity("entity1", "foo"), + new TestEntity("entity1", "bar"))))) + .hasMessageThat() + .contains("Multiple entities of type TestEntity with the same ID"); + } + + @Test + void putAll_deepCopiesNewEntitiesWithoutExtraSelects() { + persistResource(theEntity); + TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated"); + TestEntity newEntity1 = new TestEntity("newEntity1", "data1"); + TestEntity newEntity2 = new TestEntity("newEntity2", "data2"); + tm().transact( + () -> { + Statistics stats = + tm().getEntityManager() + .getEntityManagerFactory() + .unwrap(SessionFactory.class) + .getStatistics(); + stats.setStatisticsEnabled(true); + stats.clear(); + tm().putAll(updatedTheEntity, newEntity1, newEntity2); + // Only 1 SQL statement (the single batched findMultiple SELECT) should have been + // prepared so far; merging newEntity1 and newEntity2 must not trigger extra SELECTs. + assertThat(stats.getPrepareStatementCount()).isEqualTo(1); + assertThat(tm().getEntityManager().contains(updatedTheEntity)).isFalse(); + assertThat(tm().getEntityManager().contains(newEntity1)).isFalse(); + assertThat(tm().getEntityManager().contains(newEntity2)).isFalse(); + newEntity1.data = "mutated_in_memory"; + }); + assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity); + assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity1")))) + .isEqualTo(new TestEntity("newEntity1", "data1")); + assertThat(tm().transact(() -> tm().loadByKey(VKey.create(TestEntity.class, "newEntity2")))) + .isEqualTo(new TestEntity("newEntity2", "data2")); + } + @Test void update_succeeds() { persistResource(theEntity); @@ -518,11 +590,54 @@ void updateAll_succeeds() { new TestEntity("entity1", "foo_updated"), new TestEntity("entity2", "bar_updated"), new TestEntity("entity3", "qux_updated")); - tm().transact(() -> tm().updateAll(updated)); + tm().transact( + () -> { + Statistics stats = + tm().getEntityManager() + .getEntityManagerFactory() + .unwrap(SessionFactory.class) + .getStatistics(); + stats.setStatisticsEnabled(true); + stats.clear(); + tm().updateAll(updated); + // Only 1 SQL statement (the single batched findMultiple SELECT) should have been + // prepared before flush; merging the entities into the pre-warmed persistence context + // must not trigger extra per-entity SELECTs. + assertThat(stats.getPrepareStatementCount()).isEqualTo(1); + }); assertThat(tm().transact(() -> tm().loadAllOf(TestEntity.class))) .containsExactlyElementsIn(updated); } + @Test + void updateAll_mixedEntityTypes_succeeds() { + persistResource(theEntity); + persistResource(compoundIdEntity); + TestEntity updatedTheEntity = new TestEntity("theEntity", "foo_updated"); + TestCompoundIdEntity updatedCompoundEntity = + new TestCompoundIdEntity("compoundIdEntity", 10, "bar_updated"); + tm().transact(() -> tm().updateAll(updatedTheEntity, updatedCompoundEntity)); + assertThat(tm().transact(() -> tm().loadByKey(theEntityKey))).isEqualTo(updatedTheEntity); + assertThat(tm().transact(() -> tm().loadByKey(compoundIdEntityKey))) + .isEqualTo(updatedCompoundEntity); + } + + @Test + void updateAll_duplicateIds_throws() { + persistResource(theEntity); + assertThat( + assertThrows( + IllegalArgumentException.class, + () -> + tm().transact( + () -> + tm().updateAll( + new TestEntity("theEntity", "foo_updated"), + new TestEntity("theEntity", "bar_updated"))))) + .hasMessageThat() + .contains("Multiple entities of type TestEntity with the same ID"); + } + @Test void updateAll_rollsBackWhenFailure() { persistResources(moreEntities);