From 8b5920ca35c65a363d94ec9c3ab605c0f1c5e559 Mon Sep 17 00:00:00 2001
From: Andrea Boriero <andrea@hibernate.org>
Date: Fri, 24 Jan 2025 16:16:19 +0100
Subject: [PATCH] HHH-19083 Allow Hibernate Reactive to register a
 org.hibernate.loader.ast.spi.AfterLoadAction

---
 .../internal/StatefulPersistenceContext.java  |  7 ++++++-
 .../hibernate/engine/spi/EntityHolder.java    |  8 +++++++
 .../internal/DeferredResultSetAccess.java     | 21 ++++++++++++-------
 ...luesSourceProcessingStateStandardImpl.java |  6 ++++--
 4 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java b/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
index 79a45e274b41..7eea08b82e7c 100644
--- a/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
+++ b/hibernate-core/src/main/java/org/hibernate/engine/internal/StatefulPersistenceContext.java
@@ -489,7 +489,7 @@ private void processLoadedEntityHolder(
 						getSession()
 				);
 			}
-			( (EntityHolderImpl) holder ).entityInitializer = null;
+			holder.resetEntityInitialier();
 		}
 	}
 
@@ -2297,6 +2297,11 @@ public boolean isDetached() {
 			return state == EntityHolderState.DETACHED;
 		}
 
+		@Override
+		public void resetEntityInitialier(){
+			entityInitializer = null;
+		}
+
 		public EntityHolderImpl withEntity(EntityKey entityKey, EntityPersister descriptor, Object entity) {
 			return withData( entityKey, descriptor, entity, null );
 		}
diff --git a/hibernate-core/src/main/java/org/hibernate/engine/spi/EntityHolder.java b/hibernate-core/src/main/java/org/hibernate/engine/spi/EntityHolder.java
index 1f1018d36508..7695867f0590 100644
--- a/hibernate-core/src/main/java/org/hibernate/engine/spi/EntityHolder.java
+++ b/hibernate-core/src/main/java/org/hibernate/engine/spi/EntityHolder.java
@@ -70,4 +70,12 @@ public interface EntityHolder {
 	 * Whether the entity is detached.
 	 */
 	boolean isDetached();
+
+
+	/**
+	 * For Hibernate Reactive
+	 *
+	 * Set the EntityInitializer to null
+	 */
+	void resetEntityInitialier();
 }
diff --git a/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/DeferredResultSetAccess.java b/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/DeferredResultSetAccess.java
index 6ccda046871c..e467e7aa2d8d 100644
--- a/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/DeferredResultSetAccess.java
+++ b/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/DeferredResultSetAccess.java
@@ -112,13 +112,7 @@ public DeferredResultSetAccess(
 						lockOptionsToUse.setTimeOut( lockOptions.getTimeOut() );
 						lockOptionsToUse.setLockScope( lockOptions.getLockScope() );
 
-						executionContext.getCallback().registerAfterLoadAction( (entity, persister, session) ->
-								session.asSessionImplementor().lock(
-										persister.getEntityName(),
-										entity,
-										lockOptionsToUse
-								)
-						);
+						registerAfterLoadAction( executionContext, lockOptionsToUse );
 					}
 				}
 				else {
@@ -136,6 +130,19 @@ public DeferredResultSetAccess(
 		}
 	}
 
+	/**
+	 * For Hibernate Reactive
+	 */
+	protected void registerAfterLoadAction(ExecutionContext executionContext, LockOptions lockOptionsToUse) {
+		executionContext.getCallback().registerAfterLoadAction( (entity, persister, session) ->
+				session.asSessionImplementor().lock(
+						persister.getEntityName(),
+						entity,
+						lockOptionsToUse
+				)
+		);
+	}
+
 	private static boolean useFollowOnLocking(
 			JdbcLockStrategy jdbcLockStrategy,
 			String sql,
diff --git a/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java b/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java
index f666d3302ac7..0faf7f289387 100644
--- a/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java
+++ b/hibernate-core/src/main/java/org/hibernate/sql/results/jdbc/internal/JdbcValuesSourceProcessingStateStandardImpl.java
@@ -145,8 +145,10 @@ else if ( getSession() instanceof EventSource ) {
 		}
 	}
 
-
-	private void finishLoadingCollections() {
+	/**
+	 * For Hibernate Reactive
+	 */
+	public void finishLoadingCollections() {
 		if ( loadingCollectionMap != null ) {
 			for ( LoadingCollectionEntry loadingCollectionEntry : loadingCollectionMap.values() ) {
 				loadingCollectionEntry.finishLoading( getExecutionContext() );