HHH-9937 - bytecode enhancement - expose loaded state in LazyAttributeLoader

Author: barreiro

hibernate-core/src/main/java/org/hibernate/Hibernate.java

@@ -8,12 +8,15 @@
 
 import java.util.Iterator;
 
+import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributeLoader;
 import org.hibernate.bytecode.instrumentation.internal.FieldInterceptionHelper;
 import org.hibernate.bytecode.instrumentation.spi.FieldInterceptor;
 import org.hibernate.collection.spi.PersistentCollection;
 import org.hibernate.engine.HibernateIterator;
 import org.hibernate.engine.jdbc.LobCreator;
 import org.hibernate.engine.jdbc.spi.JdbcServices;
+import org.hibernate.engine.spi.PersistentAttributeInterceptable;
+import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.proxy.HibernateProxy;
 import org.hibernate.proxy.LazyInitializer;
@@ -169,13 +172,19 @@ public static boolean isPropertyInitialized(Object proxy, String propertyName) {
 			entity = proxy;
 		}
 
+		if ( entity instanceof PersistentAttributeInterceptable ) {
+			PersistentAttributeInterceptor interceptor = ( (PersistentAttributeInterceptable) entity ).$$_hibernate_getInterceptor();
+			if ( interceptor != null && interceptor instanceof LazyAttributeLoader ) {
+				return ( (LazyAttributeLoader) interceptor ).isAttributeLoaded( propertyName );
+			}
+		}
+
 		if ( FieldInterceptionHelper.isInstrumented( entity ) ) {
 			final FieldInterceptor interceptor = FieldInterceptionHelper.extractFieldInterceptor( entity );
 			return interceptor == null || interceptor.isInitialized( propertyName );
 		}
-		else {
-			return true;
-		}
+
+		return true;
 	}
 
 }

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/spi/interceptor/LazyAttributeLoader.java

@@ -22,7 +22,7 @@
  */
 public class LazyAttributeLoader implements PersistentAttributeInterceptor {
 
-	private final transient SessionImplementor session;
+	private transient SessionImplementor session;
 	private final Set<String> lazyFields;
 	private final String entityName;
 
@@ -35,7 +35,7 @@ public LazyAttributeLoader(SessionImplementor session, Set<String> lazyFields, S
 	}
 
 	protected final Object intercept(Object target, String fieldName, Object value) {
-		if ( lazyFields != null && lazyFields.contains( fieldName ) && !initializedFields.contains( fieldName ) ) {
+		if ( !isAttributeLoaded( fieldName ) ) {
 			if ( session == null ) {
 				throw new LazyInitializationException( "entity with lazy properties is not associated with a session" );
 			}
@@ -58,6 +58,25 @@ else if ( !session.isOpen() || !session.isConnected() ) {
 		}
 	}
 
+	public final void setSession(SessionImplementor session) {
+		this.session = session;
+	}
+
+	public boolean isAttributeLoaded(String fieldName) {
+		return lazyFields == null || !lazyFields.contains( fieldName ) || initializedFields.contains( fieldName );
+	}
+
+	public boolean isUninitialized() {
+		if ( lazyFields != null ) {
+			for ( String fieldName : lazyFields ) {
+				if ( !initializedFields.contains( fieldName ) ) {
+					return true;
+				}
+			}
+		}
+		return false;
+	}
+
 	public void setLoaded(String attributeName) {
 		initializedFields.add( attributeName );
 	}

hibernate-core/src/main/java/org/hibernate/persister/entity/AbstractEntityPersister.java

@@ -33,6 +33,7 @@
 import org.hibernate.Session;
 import org.hibernate.StaleObjectStateException;
 import org.hibernate.StaleStateException;
+import org.hibernate.bytecode.enhance.spi.interceptor.LazyAttributeLoader;
 import org.hibernate.bytecode.instrumentation.spi.FieldInterceptor;
 import org.hibernate.bytecode.instrumentation.spi.LazyPropertyInitializer;
 import org.hibernate.bytecode.spi.EntityInstrumentationMetadata;
@@ -62,6 +63,8 @@
 import org.hibernate.engine.spi.LoadQueryInfluencers;
 import org.hibernate.engine.spi.Mapping;
 import org.hibernate.engine.spi.PersistenceContext.NaturalIdHelper;
+import org.hibernate.engine.spi.PersistentAttributeInterceptable;
+import org.hibernate.engine.spi.PersistentAttributeInterceptor;
 import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
 import org.hibernate.engine.spi.ValueInclusion;
@@ -4196,6 +4199,13 @@ public void afterReassociate(Object entity, SessionImplementor session) {
 			}
 		}
 
+		if ( entity instanceof PersistentAttributeInterceptable ) {
+			PersistentAttributeInterceptor interceptor = ( (PersistentAttributeInterceptable) entity ).$$_hibernate_getInterceptor();
+			if ( interceptor != null && interceptor instanceof LazyAttributeLoader ) {
+				( (LazyAttributeLoader) interceptor ).setSession( session );
+			}
+		}
+
 		handleNaturalIdReattachment( entity, session );
 	}
 

hibernate-core/src/main/java/org/hibernate/tuple/entity/PojoEntityTuplizer.java

@@ -309,6 +309,12 @@ public void afterInitialize(Object entity, boolean lazyPropertiesAreUnfetched, S
 	@Override
 	public boolean hasUninitializedLazyProperties(Object entity) {
 		if ( getEntityMetamodel().hasLazyProperties() ) {
+			if ( entity instanceof PersistentAttributeInterceptable ) {
+				PersistentAttributeInterceptor interceptor = ( (PersistentAttributeInterceptable) entity ).$$_hibernate_getInterceptor();
+				if ( interceptor != null && interceptor instanceof LazyAttributeLoader ) {
+					return ( (LazyAttributeLoader) interceptor ).isUninitialized();
+				}
+			}
 			FieldInterceptor callback = FieldInterceptionHelper.extractFieldInterceptor( entity );
 			return callback != null && !callback.isInitialized();
 		}

hibernate-core/src/test/java/org/hibernate/test/bytecode/enhancement/EnhancerTest.java

@@ -6,7 +6,6 @@
  */
 package org.hibernate.test.bytecode.enhancement;
 
-import org.hibernate.test.bytecode.enhancement.lazy.LazyBasicFieldNotInitializedTestTask;
 import org.hibernate.testing.FailureExpected;
 import org.hibernate.testing.TestForIssue;
 import org.hibernate.testing.junit4.BaseUnitTestCase;
@@ -19,6 +18,7 @@
 import org.hibernate.test.bytecode.enhancement.join.HHH3949TestTask2;
 import org.hibernate.test.bytecode.enhancement.join.HHH3949TestTask3;
 import org.hibernate.test.bytecode.enhancement.join.HHH3949TestTask4;
+import org.hibernate.test.bytecode.enhancement.lazy.LazyBasicFieldNotInitializedTestTask;
 import org.hibernate.test.bytecode.enhancement.lazy.LazyLoadingIntegrationTestTask;
 import org.hibernate.test.bytecode.enhancement.lazy.LazyLoadingTestTask;
 import org.junit.Test;
@@ -75,7 +75,6 @@ public void testHHH3949() {
 
 	@Test
 	@TestForIssue( jiraKey = "HHH-9937")
-	@FailureExpected( jiraKey = "HHH-9937")
 	public void testLazyBasicFieldNotInitialized() {
 		EnhancerTestUtils.runEnhancerTestTask( LazyBasicFieldNotInitializedTestTask.class );
 	}