HHH-16952 Discover embeddable types through @Embedded annotation for …

…enhancement in a pre-discovery phase

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/ByteBuddyEnhancementContext.java

@@ -16,11 +16,13 @@
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 
 import jakarta.persistence.Embedded;
+import jakarta.persistence.metamodel.Type;
 import net.bytebuddy.description.field.FieldDescription;
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.scaffold.MethodGraph;
 import net.bytebuddy.matcher.ElementMatcher;
+import net.bytebuddy.pool.TypePool;
 
 import static net.bytebuddy.matcher.ElementMatchers.isGetter;
 
@@ -88,6 +90,50 @@ public boolean doBiDirectionalAssociationManagement(AnnotatedFieldDescription fi
 		return enhancementContext.doBiDirectionalAssociationManagement( field );
 	}
 
+	public boolean isDiscoveredType(TypeDescription typeDescription) {
+		return enhancementContext.isDiscoveredType( new UnloadedTypeDescription( typeDescription ) );
+	}
+
+	public void registerDiscoveredType(TypeDescription typeDescription, Type.PersistenceType type) {
+		enhancementContext.registerDiscoveredType( new UnloadedTypeDescription( typeDescription ), type );
+	}
+
+	public void discoverCompositeTypes(TypeDescription managedCtClass, TypePool typePool) {
+		if ( isDiscoveredType( managedCtClass ) ) {
+			return;
+		}
+		final Type.PersistenceType determinedPersistenceType;
+		if ( isEntityClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.ENTITY;
+		}
+		else if ( isCompositeClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.EMBEDDABLE;
+		}
+		else if ( isMappedSuperclassClass( managedCtClass ) ) {
+			determinedPersistenceType = Type.PersistenceType.MAPPED_SUPERCLASS;
+		}
+		else {
+			// Default to assuming a basic type if this is not a managed type
+			determinedPersistenceType = Type.PersistenceType.BASIC;
+		}
+		registerDiscoveredType( managedCtClass, determinedPersistenceType );
+		if ( determinedPersistenceType != Type.PersistenceType.BASIC ) {
+			final EnhancerImpl.AnnotatedFieldDescription[] enhancedFields = PersistentAttributeTransformer.collectPersistentFields(
+							managedCtClass,
+							this,
+							typePool
+					)
+					.getEnhancedFields();
+			for ( EnhancerImpl.AnnotatedFieldDescription enhancedField : enhancedFields ) {
+				final TypeDescription type = enhancedField.getType().asErasure();
+				if ( !type.isInterface() && enhancedField.hasAnnotation( Embedded.class ) ) {
+					registerDiscoveredType( type, Type.PersistenceType.EMBEDDABLE );
+				}
+				discoverCompositeTypes( type, typePool );
+			}
+		}
+	}
+
 	Optional<MethodDescription> resolveGetter(FieldDescription fieldDescription) {
 		Map<String, MethodDescription> getters = getterByTypeMap
 				.computeIfAbsent( fieldDescription.getDeclaringType().asErasure(), declaringType -> {

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/EnhancerImpl.java

@@ -46,6 +46,7 @@
 import jakarta.persistence.Access;
 import jakarta.persistence.AccessType;
 import jakarta.persistence.Transient;
+import jakarta.persistence.metamodel.Type;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.annotation.AnnotationDescription;
 import net.bytebuddy.description.annotation.AnnotationList;
@@ -161,6 +162,16 @@ public byte[] enhance(String className, byte[] originalBytes) throws Enhancement
 		}
 	}
 
+	@Override
+	public void discoverTypes(String className, byte[] originalBytes) {
+		if ( originalBytes != null ) {
+			classFileLocator.setClassNameAndBytes( className, originalBytes );
+		}
+		final TypeDescription typeDescription = typePool.describe( className ).resolve();
+		enhancementContext.registerDiscoveredType( typeDescription, Type.PersistenceType.ENTITY );
+		enhancementContext.discoverCompositeTypes( typeDescription, typePool );
+	}
+
 	private TypePool buildTypePool(final ClassFileLocator classFileLocator) {
 		return TypePool.Default.WithLazyResolution.of( classFileLocator );
 	}

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/FieldAccessEnhancer.java

@@ -68,6 +68,8 @@ public void visitFieldInsn(int opcode, String owner, String name, String desc) {
 
 				TypeDescription declaredOwnerType = findDeclaredType( owner );
 				AnnotatedFieldDescription field = findField( declaredOwnerType, name, desc );
+				// try to discover composite types on the fly to support some testing scenarios
+				enhancementContext.discoverCompositeTypes( declaredOwnerType, typePool );
 
 				if ( ( enhancementContext.isEntityClass( declaredOwnerType.asErasure() )
 						|| enhancementContext.isCompositeClass( declaredOwnerType.asErasure() ) )

hibernate-core/src/main/java/org/hibernate/bytecode/enhance/internal/bytebuddy/PersistentAttributeTransformer.java

@@ -103,6 +103,10 @@ private PersistentAttributeTransformer(
 		this.enhancedFields = enhancedFields;
 	}
 
+	public AnnotatedFieldDescription[] getEnhancedFields() {
+		return enhancedFields;
+	}
+
 	public static PersistentAttributeTransformer collectPersistentFields(
 			TypeDescription managedCtClass,
 			ByteBuddyEnhancementContext enhancementContext,

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

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import java.util.concurrent.ConcurrentHashMap;
+
 import jakarta.persistence.Basic;
 import jakarta.persistence.Convert;
 import jakarta.persistence.ElementCollection;
@@ -15,6 +17,7 @@
 import jakarta.persistence.MappedSuperclass;
 import jakarta.persistence.OneToMany;
 import jakarta.persistence.Transient;
+import jakarta.persistence.metamodel.Type;
 
 /**
  * default implementation of EnhancementContext. May be sub-classed as needed.
@@ -23,6 +26,8 @@
  */
 public class DefaultEnhancementContext implements EnhancementContext {
 
+	private final ConcurrentHashMap<String, Type.PersistenceType> discoveredTypes = new ConcurrentHashMap<>();
+
 	/**
 	 * @return the classloader for this class
 	 */
@@ -44,7 +49,8 @@ public boolean isEntityClass(UnloadedClass classDescriptor) {
 	 */
 	@Override
 	public boolean isCompositeClass(UnloadedClass classDescriptor) {
-		return classDescriptor.hasAnnotation( Embeddable.class );
+		return classDescriptor.hasAnnotation( Embeddable.class )
+				|| discoveredTypes.get( classDescriptor.getName() ) == Type.PersistenceType.EMBEDDABLE;
 	}
 
 	/**
@@ -124,4 +130,14 @@ public boolean isMappedCollection(UnloadedField field) {
 	public UnloadedField[] order(UnloadedField[] persistentFields) {
 		return persistentFields;
 	}
+
+	@Override
+	public boolean isDiscoveredType(UnloadedClass classDescriptor) {
+		return discoveredTypes.containsKey( classDescriptor.getName() );
+	}
+
+	@Override
+	public void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type) {
+		discoveredTypes.put( classDescriptor.getName(), type );
+	}
 }

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

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import jakarta.persistence.metamodel.Type;
+
 /**
  * The context for performing an enhancement.  Enhancement can happen in any number of ways:<ul>
  * <li>Build time, via Ant</li>
@@ -140,4 +142,8 @@ public interface EnhancementContext {
 	 * @return {@code true} if the field is mapped
 	 */
 	boolean isMappedCollection(UnloadedField field);
+
+	boolean isDiscoveredType(UnloadedClass classDescriptor);
+
+	void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type);
 }

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

@@ -6,6 +6,8 @@
  */
 package org.hibernate.bytecode.enhance.spi;
 
+import jakarta.persistence.metamodel.Type;
+
 public class EnhancementContextWrapper implements EnhancementContext {
 
 	private final ClassLoader loadingClassloader;
@@ -75,4 +77,14 @@ public boolean isLazyLoadable(UnloadedField field) {
 	public boolean isMappedCollection(UnloadedField field) {
 		return wrappedContext.isMappedCollection( field );
 	}
+
+	@Override
+	public boolean isDiscoveredType(UnloadedClass classDescriptor) {
+		return wrappedContext.isDiscoveredType( classDescriptor );
+	}
+
+	@Override
+	public void registerDiscoveredType(UnloadedClass classDescriptor, Type.PersistenceType type) {
+		wrappedContext.registerDiscoveredType( classDescriptor, type );
+	}
 }

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

@@ -30,4 +30,6 @@ public interface Enhancer {
 	 * @throws EnhancementException Indicates a problem performing the enhancement
 	 */
 	byte[] enhance(String className, byte[] originalBytes) throws EnhancementException;
+
+	void discoverTypes(String className, byte[] originalBytes);
 }

hibernate-core/src/main/java/org/hibernate/bytecode/spi/ClassTransformer.java

@@ -42,4 +42,6 @@ byte[] transform(
 			@Nullable Class<?> classBeingRedefined,
 			ProtectionDomain protectionDomain,
 			byte[] classfileBuffer)  throws TransformerException;
+
+	void discoverTypes(ClassLoader loader, String entityClassName);
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/EntityManagerFactoryBuilderImpl.java

@@ -52,6 +52,7 @@
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
 import org.hibernate.bytecode.enhance.spi.UnloadedClass;
 import org.hibernate.bytecode.enhance.spi.UnloadedField;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.cfg.Environment;
 import org.hibernate.engine.jdbc.connections.spi.ConnectionProvider;
@@ -70,6 +71,7 @@
 import org.hibernate.jpa.boot.spi.TypeContributorList;
 import org.hibernate.jpa.internal.util.LogHelper;
 import org.hibernate.jpa.internal.util.PersistenceUnitTransactionTypeHelper;
+import org.hibernate.mapping.PersistentClass;
 import org.hibernate.proxy.EntityNotFoundDelegate;
 import org.hibernate.resource.transaction.backend.jdbc.internal.JdbcResourceLocalTransactionCoordinatorBuilderImpl;
 import org.hibernate.resource.transaction.backend.jta.internal.JtaTransactionCoordinatorBuilderImpl;
@@ -339,6 +341,20 @@ private EntityManagerFactoryBuilderImpl(
 				);
 
 				persistenceUnit.pushClassTransformer( enhancementContext );
+
+				if ( !persistenceUnit.getClassTransformers().isEmpty() ) {
+					final ClassLoader classLoader = persistenceUnit.getTempClassLoader();
+					if ( classLoader == null ) {
+						throw persistenceException( "Enhancement requires a temp class loader, but none was given." );
+					}
+					for ( ClassTransformer classTransformer : persistenceUnit.getClassTransformers() ) {
+						for ( PersistentClass entityBinding : metadata.getEntityBindings() ) {
+							if ( entityBinding.getClassName() != null ) {
+								classTransformer.discoverTypes( classLoader, entityBinding.getClassName() );
+							}
+						}
+					}
+				}
 			}
 
 			// for the time being we want to revoke access to the temp ClassLoader if one was passed

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/ParsedPersistenceXmlDescriptor.java

@@ -9,13 +9,16 @@
 import java.net.URL;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
 import jakarta.persistence.ValidationMode;
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 
 /**
  * Describes the information gleaned from a {@code <persistence-unit/>} element in a {@code persistence.xml} file
@@ -189,4 +192,9 @@ public ClassLoader getTempClassLoader() {
 	public void pushClassTransformer(EnhancementContext enhancementContext) {
 		// todo : log a message that this is currently not supported...
 	}
+
+	@Override
+	public Collection<ClassTransformer> getClassTransformers() {
+		return Collections.emptyList();
+	}
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/internal/PersistenceUnitInfoDescriptor.java

@@ -7,6 +7,8 @@
 package org.hibernate.jpa.boot.internal;
 
 import java.net.URL;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
@@ -15,6 +17,7 @@
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 import org.hibernate.jpa.boot.spi.PersistenceUnitDescriptor;
 import org.hibernate.jpa.internal.enhance.EnhancingClassTransformerImpl;
 
@@ -23,6 +26,7 @@
  */
 public class PersistenceUnitInfoDescriptor implements PersistenceUnitDescriptor {
 	private final PersistenceUnitInfo persistenceUnitInfo;
+	private final ArrayList<ClassTransformer> classTransformers = new ArrayList<>();
 
 	public PersistenceUnitInfoDescriptor(PersistenceUnitInfo persistenceUnitInfo) {
 		this.persistenceUnitInfo = persistenceUnitInfo;
@@ -110,6 +114,16 @@ public List<URL> getJarFileUrls() {
 
 	@Override
 	public void pushClassTransformer(EnhancementContext enhancementContext) {
-		persistenceUnitInfo.addTransformer( new EnhancingClassTransformerImpl( enhancementContext ) );
+		// During testing, we will return a null temp class loader in cases where we don't care about enhancement
+		if ( persistenceUnitInfo.getNewTempClassLoader() != null ) {
+			final EnhancingClassTransformerImpl classTransformer = new EnhancingClassTransformerImpl( enhancementContext );
+			classTransformers.add( classTransformer );
+			persistenceUnitInfo.addTransformer( classTransformer );
+		}
+	}
+
+	@Override
+	public Collection<ClassTransformer> getClassTransformers() {
+		return classTransformers;
 	}
 }

hibernate-core/src/main/java/org/hibernate/jpa/boot/spi/PersistenceUnitDescriptor.java

@@ -7,13 +7,15 @@
 package org.hibernate.jpa.boot.spi;
 
 import java.net.URL;
+import java.util.Collection;
 import java.util.List;
 import java.util.Properties;
 import jakarta.persistence.SharedCacheMode;
 import jakarta.persistence.ValidationMode;
 import jakarta.persistence.spi.PersistenceUnitTransactionType;
 
 import org.hibernate.bytecode.enhance.spi.EnhancementContext;
+import org.hibernate.bytecode.spi.ClassTransformer;
 
 /**
  * Abstraction for dealing with either {@code <persistence-unit/>} information whether that comes from
@@ -85,4 +87,6 @@ public interface PersistenceUnitDescriptor {
 	ClassLoader getTempClassLoader();
 
 	void pushClassTransformer(EnhancementContext enhancementContext);
+
+	Collection<ClassTransformer> getClassTransformers();
 }