From e96741cfa6d9b0d82b74020734321624daacb4ec Mon Sep 17 00:00:00 2001
From: Yanming Zhou <zhouyanming@gmail.com>
Date: Fri, 24 May 2024 10:32:27 +0800
Subject: [PATCH] HHH-18164 Make @IdGeneratorType respect BeanContainer

---
 .../boot/model/internal/GeneratorBinder.java  | 57 ++++++++++++++++---
 .../boot/model/internal/PropertyBinder.java   |  5 ++
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/hibernate-core/src/main/java/org/hibernate/boot/model/internal/GeneratorBinder.java b/hibernate-core/src/main/java/org/hibernate/boot/model/internal/GeneratorBinder.java
index 6c26babf48..2a1cfc183b 100644
--- a/hibernate-core/src/main/java/org/hibernate/boot/model/internal/GeneratorBinder.java
+++ b/hibernate-core/src/main/java/org/hibernate/boot/model/internal/GeneratorBinder.java
@@ -65,8 +65,8 @@ import org.hibernate.models.spi.ClassDetails;
 import org.hibernate.models.spi.MemberDetails;
 import org.hibernate.mapping.Value;
 import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.container.spi.ContainedBean;
 import org.hibernate.resource.beans.spi.BeanInstanceProducer;
-import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
 
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.type.Type;
@@ -552,14 +552,53 @@ public class GeneratorBinder {
 		final Class<? extends Generator> generatorClass = idGeneratorType.value();
 		return creationContext -> {
 			checkGeneratorClass( generatorClass );
-			final Generator generator = instantiateGenerator(
-					annotation,
-					beanContainer,
-					creationContext,
-					generatorClass,
-					idAttributeMember,
-					annotationType
-			);
+			Generator generator;
+			if ( beanContainer != null ) {
+				final ContainedBean<? extends Generator> bean = beanContainer.getBean(
+						generatorClass,
+						new BeanContainer.LifecycleOptions() {
+							@Override
+							public boolean canUseCachedReferences() {
+								return false;
+							}
+
+							@Override
+							public boolean useJpaCompliantCreation() {
+								return true;
+							}
+						},
+						new BeanInstanceProducer() {
+							@SuppressWarnings("unchecked")
+							@Override
+							public <B> B produceBeanInstance(Class<B> beanType) {
+								return (B) instantiateGenerator(
+										annotation,
+										idAttributeMember,
+										annotationType,
+										CustomIdGeneratorCreationContext.class,
+										generatorClass,
+										creationContext
+								);
+							}
+
+							@Override
+							public <B> B produceBeanInstance(String name, Class<B> beanType) {
+								return produceBeanInstance( beanType );
+							}
+						}
+				);
+				generator = bean.getBeanInstance();
+			}
+			else {
+				generator = instantiateGenerator(
+						annotation,
+						beanContainer,
+						creationContext,
+						generatorClass,
+						idAttributeMember,
+						annotationType
+				);
+			}
 			callInitialize( annotation, idAttributeMember, creationContext, generator );
 			callConfigure( creationContext, generator );
 			checkIdGeneratorTiming( annotationType, generator );
diff --git a/hibernate-core/src/main/java/org/hibernate/boot/model/internal/PropertyBinder.java b/hibernate-core/src/main/java/org/hibernate/boot/model/internal/PropertyBinder.java
index 8173e4c63d..47996b983e 100644
--- a/hibernate-core/src/main/java/org/hibernate/boot/model/internal/PropertyBinder.java
+++ b/hibernate-core/src/main/java/org/hibernate/boot/model/internal/PropertyBinder.java
@@ -63,6 +63,11 @@ import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
 import org.hibernate.service.ServiceRegistry;
 import org.hibernate.usertype.CompositeUserType;
 
+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;
 
 import jakarta.persistence.Basic;