HHH-18164 Make @IdGeneratorType respect BeanContainer

hibernate-core/src/main/java/org/hibernate/boot/model/internal/GeneratorBinder.java

@@ -43,6 +43,8 @@ import org.hibernate.mapping.GeneratorCreator;
 import org.hibernate.mapping.IdentifierGeneratorCreator;
 import org.hibernate.mapping.SimpleValue;
 import org.hibernate.mapping.Value;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
 
 import org.jboss.logging.Logger;
 
@@ -373,7 +375,7 @@ public class GeneratorBinder {
 		};
 	}
 
-	static IdentifierGeneratorCreator identifierGeneratorCreator(XProperty idProperty, Annotation annotation) {
+	static IdentifierGeneratorCreator identifierGeneratorCreator(XProperty idProperty, Annotation annotation, BeanContainer beanContainer) {
 		final Member member = HCANNHelper.getUnderlyingMember( idProperty );
 		final Class<? extends Annotation> annotationType = annotation.annotationType();
 		final IdGeneratorType idGeneratorType = annotationType.getAnnotation( IdGeneratorType.class );
@@ -381,15 +383,48 @@ public class GeneratorBinder {
 		return creationContext -> {
 			final Class<? extends Generator> generatorClass = idGeneratorType.value();
 			checkGeneratorClass( generatorClass );
-			final Generator generator =
+			Generator generator;
-					instantiateGenerator(
+			if ( beanContainer != null ) {
-							annotation,
+				generator = beanContainer.getBean( generatorClass, new BeanContainer.LifecycleOptions() {
-							member,
+					@Override
-							annotationType,
+					public boolean canUseCachedReferences() {
-							creationContext,
+						return false;
-							CustomIdGeneratorCreationContext.class,
+					}
-							generatorClass
+
-					);
+					@Override
+					public boolean useJpaCompliantCreation() {
+						return true;
+					}
+				}, new BeanInstanceProducer() {
+					@SuppressWarnings( "unchecked" )
+					@Override
+					public <B> B produceBeanInstance(Class<B> beanType) {
+						return (B) instantiateGenerator(
+								annotation,
+								member,
+								annotationType,
+								creationContext,
+								CustomIdGeneratorCreationContext.class,
+								generatorClass
+						);
+					}
+
+					@Override
+					public <B> B produceBeanInstance(String name, Class<B> beanType) {
+						return produceBeanInstance(beanType);
+					}
+				} ).getBeanInstance();
+			}
+			else {
+				generator = instantiateGenerator(
+						annotation,
+						member,
+						annotationType,
+						creationContext,
+						CustomIdGeneratorCreationContext.class,
+						generatorClass
+				);
+			}
 			callInitialize( annotation, member, creationContext, generator );
 			callConfigure( creationContext, generator );
 			checkIdGeneratorTiming( annotationType, generator );

hibernate-core/src/main/java/org/hibernate/boot/model/internal/PropertyBinder.java

@@ -69,6 +69,10 @@ import org.hibernate.mapping.ToOne;
 import org.hibernate.mapping.Value;
 import org.hibernate.metamodel.spi.EmbeddableInstantiator;
 
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.internal.Helper;
+import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
+import org.hibernate.service.ServiceRegistry;
 import org.hibernate.usertype.CompositeUserType;
 import org.jboss.logging.Logger;
 
@@ -1430,7 +1434,9 @@ public class PropertyBinder {
 					+ "' has too many generator annotations " + combine( idGeneratorAnnotations, generatorAnnotations ) );
 		}
 		if ( !idGeneratorAnnotations.isEmpty() ) {
-			idValue.setCustomIdGeneratorCreator( identifierGeneratorCreator( idProperty, idGeneratorAnnotations.get(0) ) );
+			final ServiceRegistry serviceRegistry = context.getBootstrapContext().getServiceRegistry();
+			final BeanContainer beanContainer = Helper.allowExtensionsInCdi( serviceRegistry ) ? serviceRegistry.requireService( ManagedBeanRegistry.class ).getBeanContainer() : null;
+			idValue.setCustomIdGeneratorCreator( identifierGeneratorCreator( idProperty, idGeneratorAnnotations.get(0), beanContainer ) );
 		}
 		else if ( !generatorAnnotations.isEmpty() ) {
 //			idValue.setCustomGeneratorCreator( generatorCreator( idProperty, generatorAnnotation ) );

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/IdGeneratorTypeWithBeanContainerTest.java

@@ -0,0 +1,46 @@
+package org.hibernate.orm.test.idgen.userdefined;
+
+import org.hibernate.cfg.AvailableSettings;
+
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.JiraKey;
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.Test;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author Yanming Zhou
+ */
+@JiraKey( "HHH-18164" )
+@SessionFactory
+@ServiceRegistry(
+		settings = {
+			@Setting(name = AvailableSettings.ALLOW_EXTENSIONS_IN_CDI, value = "true"),
+			@Setting(name = AvailableSettings.BEAN_CONTAINER, value = "org.hibernate.orm.test.idgen.userdefined.SimpleBeanContainer")
+		}
+)
+@DomainModel(annotatedClasses = IdGeneratorTypeWithBeanContainerTest.SimpleEntity.class)
+public class IdGeneratorTypeWithBeanContainerTest {
+
+	@Test void test(SessionFactoryScope scope) {
+		SimpleEntity entity = new SimpleEntity();
+		scope.inTransaction(s -> s.persist(entity));
+		assertThat(entity.id, is(SimpleBeanContainer.INITIAL_VALUE));
+	}
+
+	@Entity
+	public static class SimpleEntity {
+		@Id @SimpleId
+		long id;
+		String data;
+	}
+
+}

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/SimpleBeanContainer.java

@@ -0,0 +1,41 @@
+package org.hibernate.orm.test.idgen.userdefined;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.container.spi.ContainedBean;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+
+/**
+ *
+ */
+@SuppressWarnings("unchecked")
+public class SimpleBeanContainer implements BeanContainer {
+
+	public static final long INITIAL_VALUE = 23L;
+
+	@Override
+	public <B> ContainedBean<B> getBean(
+			Class<B> beanType,
+			LifecycleOptions lifecycleOptions,
+			BeanInstanceProducer fallbackProducer) {
+		if ( beanType == SimpleGenerator.class ) {
+			return () -> (B) new SimpleGenerator( new AtomicLong( INITIAL_VALUE ) );
+		}
+		return null;
+	}
+
+	@Override
+	public <B> ContainedBean<B> getBean(
+			String name,
+			Class<B> beanType,
+			LifecycleOptions lifecycleOptions,
+			BeanInstanceProducer fallbackProducer) {
+		return null;
+	}
+
+	@Override
+	public void stop() {
+
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/SimpleGenerator.java

@@ -0,0 +1,34 @@
+package org.hibernate.orm.test.idgen.userdefined;
+
+import java.util.EnumSet;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.hibernate.engine.spi.SharedSessionContractImplementor;
+import org.hibernate.generator.BeforeExecutionGenerator;
+import org.hibernate.generator.EventType;
+
+/**
+ * @author Yanming Zhou
+ */
+public class SimpleGenerator implements BeforeExecutionGenerator {
+
+    private final AtomicLong sequence;
+
+    public SimpleGenerator(AtomicLong sequence) {
+        this.sequence = sequence;
+    }
+
+    @Override
+    public Object generate(
+            SharedSessionContractImplementor session,
+            Object owner,
+            Object currentValue,
+            EventType eventType) {
+        return sequence.getAndIncrement();
+    }
+
+    @Override
+    public EnumSet<EventType> getEventTypes() {
+        return EnumSet.of( EventType.INSERT );
+    }
+}

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/SimpleId.java

@@ -0,0 +1,18 @@
+package org.hibernate.orm.test.idgen.userdefined;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+import org.hibernate.annotations.IdGeneratorType;
+
+/**
+ * @author Yanming Zhou
+ */
+
+@Target(ElementType.FIELD)
+@Retention(RetentionPolicy.RUNTIME)
+@IdGeneratorType(SimpleGenerator.class)
+public @interface SimpleId {
+}