2727import org .jspecify .annotations .Nullable ;
2828
2929import org .springframework .beans .BeanUtils ;
30+ import org .springframework .beans .factory .config .RuntimeBeanReference ;
3031import org .springframework .core .annotation .AnnotatedElementUtils ;
3132import org .springframework .core .annotation .MergedAnnotation ;
3233import org .springframework .core .annotation .MergedAnnotations ;
@@ -120,7 +121,16 @@ protected void customizeClass(AotRepositoryClassBuilder classBuilder) {
120121 @ Override
121122 protected void customizeConstructor (AotRepositoryConstructorBuilder constructorBuilder ) {
122123
123- constructorBuilder .addParameter ("entityManager" , EntityManager .class );
124+ String entityManagerFactoryRef = getEntityManagerFactoryRef ();
125+
126+ constructorBuilder .addParameter ("entityManager" , EntityManager .class , customizer -> {
127+
128+ customizer .bindToField ().origin (
129+ StringUtils .hasText (entityManagerFactoryRef )
130+ ? new RuntimeBeanReference (entityManagerFactoryRef , EntityManager .class )
131+ : new RuntimeBeanReference (EntityManager .class ));
132+ });
133+
124134 constructorBuilder .addParameter ("context" , RepositoryFactoryBeanSupport .FragmentCreationContext .class );
125135
126136 Optional <Class <QueryEnhancerSelector >> queryEnhancerSelector = getQueryEnhancerSelectorClass ();
@@ -135,6 +145,11 @@ protected void customizeConstructor(AotRepositoryConstructorBuilder constructorB
135145 });
136146 }
137147
148+ private String getEntityManagerFactoryRef () {
149+ return context .getConfigurationSource ().getAttribute ("entityManagerFactoryRef" )
150+ .filter (it -> !"entityManagerFactory" .equals (it )).orElse (null );
151+ }
152+
138153 @ SuppressWarnings ({ "rawtypes" , "unchecked" })
139154 private Optional <Class <QueryEnhancerSelector >> getQueryEnhancerSelectorClass () {
140155 return (Optional ) context .getConfigurationSource ().getAttribute ("queryEnhancerSelector" , Class .class )
0 commit comments