HHH-Allow Hibernate Reactive to register a org.hibernate.loader.ast.spi.AfterLoadAction

Author: dreab8

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 );
 		}

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();
 }

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,

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() );