[#1504] Added forced loading of proxy entity references

Author: blafond

hibernate-reactive-core/src/main/java/org/hibernate/reactive/engine/impl/EntityTypes.java

@@ -6,28 +6,38 @@
 package org.hibernate.reactive.engine.impl;
 
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 
 import org.hibernate.AssertionFailure;
 import org.hibernate.HibernateException;
+import org.hibernate.bytecode.enhance.spi.interceptor.EnhancementAsProxyLazinessInterceptor;
 import org.hibernate.engine.spi.EntityKey;
 import org.hibernate.engine.spi.EntityUniqueKey;
 import org.hibernate.engine.spi.PersistenceContext;
+import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.SharedSessionContractImplementor;
 import org.hibernate.persister.entity.EntityPersister;
+import org.hibernate.proxy.HibernateProxy;
+import org.hibernate.proxy.LazyInitializer;
 import org.hibernate.reactive.persister.entity.impl.ReactiveEntityPersister;
 import org.hibernate.reactive.session.impl.ReactiveQueryExecutorLookup;
 import org.hibernate.reactive.session.impl.ReactiveSessionImpl;
+import org.hibernate.reactive.util.impl.CompletionStages;
 import org.hibernate.type.EntityType;
 import org.hibernate.type.ForeignKeyDirection;
 import org.hibernate.type.OneToOneType;
 import org.hibernate.type.Type;
 
 import static org.hibernate.bytecode.enhance.spi.LazyPropertyInitializer.UNFETCHED_PROPERTY;
+import static org.hibernate.engine.internal.ManagedTypeHelper.asPersistentAttributeInterceptable;
+import static org.hibernate.engine.internal.ManagedTypeHelper.isPersistentAttributeInterceptable;
 import static org.hibernate.property.access.internal.PropertyAccessStrategyBackRefImpl.UNKNOWN;
+import static org.hibernate.proxy.HibernateProxy.extractLazyInitializer;
 import static org.hibernate.reactive.engine.impl.ForeignKeys.getEntityIdentifierIfNotUnsaved;
+import static org.hibernate.reactive.session.impl.SessionUtil.checkEntityFound;
 import static org.hibernate.reactive.util.impl.CompletionStages.completedFuture;
 import static org.hibernate.reactive.util.impl.CompletionStages.loop;
 import static org.hibernate.reactive.util.impl.CompletionStages.nullFuture;
@@ -124,13 +134,29 @@ static CompletionStage<Object> loadByUniqueKey(
 		else {
 			return persister
 					.reactiveLoadByUniqueKey( uniqueKeyPropertyName, key, session )
-					.thenApply( loaded -> {
+					.thenApply( ukResult -> {
+						Object loadedUniqueKey = ukResult;
+						// The unique key may be a HibernateProxy object and needs to be loaded before
+						// adding to persistenceContext
+						if ( loadedUniqueKey instanceof HibernateProxy ) {
+							try {
+								CompletionStage<Object> resultObject = loadHibernateProxyEntity(
+										(HibernateProxy) loadedUniqueKey,
+										session
+								);
+								loadedUniqueKey = ((CompletableFuture)resultObject).get();
+							}
+							catch (Exception e) {
+								throw new RuntimeException( e );
+							}
+						}
+
 						// If the entity was not in the Persistence Context, but was found now,
 						// add it to the Persistence Context
-						if ( loaded != null ) {
-							persistenceContext.addEntity( euk, loaded );
+						if ( loadedUniqueKey != null ) {
+							persistenceContext.addEntity( euk, loadedUniqueKey );
 						}
-						return loaded;
+						return loadedUniqueKey;
 					} );
 
 		}
@@ -311,15 +337,28 @@ private static CompletionStage<Object> resolveIdOrUniqueKey(
 							.thenCompose( fetched -> {
 								Object idOrUniqueKey = entityType.getIdentifierOrUniqueKeyType( session.getFactory() )
 										.replace( fetched, null, session, owner, copyCache );
-								return resolve( entityType, idOrUniqueKey, owner, session );
+								Object resultObject = idOrUniqueKey;
+								if( resultObject instanceof CompletableFuture ) {
+									try {
+										resultObject = ((CompletableFuture)idOrUniqueKey).get();
+									}
+									catch (Exception e) {
+										throw new RuntimeException( e );
+									}
+								}
+
+								return resolve( entityType, resultObject, owner, session );
 							} );
 				} );
 	}
 
 	/**
 	 * see EntityType#getIdentifier(Object, SharedSessionContractImplementor)
 	 */
-	private static CompletionStage<Object> getIdentifier(EntityType entityType, Object value, SessionImplementor session) {
+	private static CompletionStage<Object> getIdentifier(
+			EntityType entityType,
+			Object value,
+			SessionImplementor session) {
 		if ( entityType.isReferenceToIdentifierProperty() ) {
 			// tolerates nulls
 			return getEntityIdentifierIfNotUnsaved( entityType.getAssociatedEntityName(), value, session );
@@ -328,17 +367,79 @@ private static CompletionStage<Object> getIdentifier(EntityType entityType, Obje
 			return nullFuture();
 		}
 
-		EntityPersister entityPersister = entityType.getAssociatedEntityPersister( session.getFactory() );
+		if( value instanceof HibernateProxy ) {
+			return getIdentifierFromHibernateProxy( entityType, (HibernateProxy)value, session );
+		}
+
+		final LazyInitializer lazyInitializer = extractLazyInitializer( value );
+		if ( lazyInitializer != null ) {
+			/*
+				If the value is a Proxy and the property access is field, the value returned by
+			 	`attributeMapping.getAttributeMetadata().getPropertyAccess().getGetter().get( object )`
+			 	is always null except for the id, we need the to use the proxy implementation to
+			 	extract the property value.
+			 */
+			value = lazyInitializer.getImplementation();
+		}
+		else if ( isPersistentAttributeInterceptable( value ) ) {
+				/*
+					If the value is an instance of PersistentAttributeInterceptable, and it is not initialized
+					we need to force initialization the get the property value
+				 */
+			final PersistentAttributeInterceptor interceptor = asPersistentAttributeInterceptable( value ).$$_hibernate_getInterceptor();
+			if ( interceptor instanceof EnhancementAsProxyLazinessInterceptor ) {
+				( (EnhancementAsProxyLazinessInterceptor) interceptor ).forceInitialize( value, null );
+			}
+		}
+		final EntityPersister entityPersister = entityType.getAssociatedEntityPersister( session.getFactory() );
 		String uniqueKeyPropertyName = entityType.getRHSUniqueKeyPropertyName();
 		Object propertyValue = entityPersister.getPropertyValue( value, uniqueKeyPropertyName );
 		// We now have the value of the property-ref we reference. However,
 		// we need to dig a little deeper, as that property might also be
 		// an entity type, in which case we need to resolve its identifier
-		Type type = entityPersister.getPropertyType( uniqueKeyPropertyName );
+		final Type type = entityPersister.getPropertyType( uniqueKeyPropertyName );
 		if ( type.isEntityType() ) {
 			propertyValue = getIdentifier( (EntityType) type, propertyValue, session );
 		}
 		return completedFuture( propertyValue );
+
+	}
+
+	private static CompletionStage<Object> getIdentifierFromHibernateProxy(EntityType entityType, HibernateProxy proxy, SharedSessionContractImplementor session) {
+		LazyInitializer initializer = proxy.getHibernateLazyInitializer();
+		final String entityName = initializer.getEntityName();
+		final Object identifier = initializer.getIdentifier();
+		return ( (ReactiveSessionImpl) session ).reactiveImmediateLoad( entityName, identifier )
+				.thenApply( entity -> {
+					checkEntityFound( session, entityName, identifier, entity );
+					initializer.setSession( session );
+					initializer.setImplementation( entity );
+					if ( entity != null ) {
+						final EntityPersister entityPersister = entityType.getAssociatedEntityPersister( session.getFactory() );
+						String uniqueKeyPropertyName = entityType.getRHSUniqueKeyPropertyName();
+						Object propertyValue = entityPersister.getPropertyValue( entity, uniqueKeyPropertyName );
+						// We now have the value of the property-ref we reference. However,
+						// we need to dig a little deeper, as that property might also be
+						// an entity type, in which case we need to resolve its identifier
+						final Type type = entityPersister.getPropertyType( uniqueKeyPropertyName );
+						if ( type.isEntityType() ) {
+							propertyValue = getIdentifier( (EntityType) type, propertyValue, (SessionImplementor) session );
+						}
+						return completedFuture( propertyValue );
+					}
+					return CompletionStages.nullFuture();
+				} );
+	}
+
+	private static CompletionStage<Object> loadHibernateProxyEntity( HibernateProxy proxy, SharedSessionContractImplementor session) {
+		LazyInitializer initializer = proxy.getHibernateLazyInitializer();
+		final String entityName = initializer.getEntityName();
+		final Object identifier = initializer.getIdentifier();
+		return ( (ReactiveSessionImpl) session ).reactiveImmediateLoad( entityName, identifier )
+				.thenApply( entity -> {
+					checkEntityFound( session, entityName, identifier, entity );
+					return entity;
+				} );
 	}
 
 }