Handle injection of several persistence units

This commit adapts the generated code so that each injection point has
a dedicated namespace in the form of a private method. That prevents
the same variable to be reused twice which lead to a compilation failure
previously.

Closes gh-30437

Author: snicoll

spring-orm/src/main/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessor.java

@@ -43,6 +43,7 @@
 import org.springframework.aot.generate.GeneratedMethod;
 import org.springframework.aot.generate.GeneratedMethods;
 import org.springframework.aot.generate.GenerationContext;
+import org.springframework.aot.generate.MethodReference.ArgumentCodeGenerator;
 import org.springframework.aot.hint.RuntimeHints;
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.PropertyValues;
@@ -769,11 +770,11 @@ private static class AotContribution implements BeanRegistrationAotContribution
 
 		private final Class<?> target;
 
-		private final Collection<InjectedElement> injectedElements;
+		private final List<InjectedElement> injectedElements;
 
 		AotContribution(Class<?> target, Collection<InjectedElement> injectedElements) {
 			this.target = target;
-			this.injectedElements = injectedElements;
+			this.injectedElements = List.copyOf(injectedElements);
 		}
 
 		@Override
@@ -797,19 +798,47 @@ public void applyTo(GenerationContext generationContext, BeanRegistrationCode be
 
 		private CodeBlock generateMethodCode(RuntimeHints hints, GeneratedClass generatedClass) {
 			CodeBlock.Builder code = CodeBlock.builder();
-			InjectionCodeGenerator injectionCodeGenerator =
-					new InjectionCodeGenerator(generatedClass.getName(), hints);
-			for (InjectedElement injectedElement : this.injectedElements) {
-				CodeBlock resourceToInject = generateResourceToInjectCode(generatedClass.getMethods(),
-						(PersistenceElement) injectedElement);
-				code.add(injectionCodeGenerator.generateInjectionCode(
-						injectedElement.getMember(), INSTANCE_PARAMETER,
-						resourceToInject));
+			if (this.injectedElements.size() == 1) {
+				code.add(generateInjectedElementMethodCode(hints, generatedClass, this.injectedElements.get(0)));
+			}
+			else {
+				for (InjectedElement injectedElement : this.injectedElements) {
+					code.addStatement(applyInjectedElement(hints, generatedClass, injectedElement));
+				}
 			}
 			code.addStatement("return $L", INSTANCE_PARAMETER);
 			return code.build();
 		}
 
+		private CodeBlock applyInjectedElement(RuntimeHints hints, GeneratedClass generatedClass, InjectedElement injectedElement) {
+			String injectedElementName = injectedElement.getMember().getName();
+			GeneratedMethod generatedMethod = generatedClass.getMethods().add(new String[] { "apply", injectedElementName }, method -> {
+				method.addJavadoc("Apply the persistence injection for '$L'.", injectedElementName);
+				method.addModifiers(javax.lang.model.element.Modifier.PRIVATE,
+						javax.lang.model.element.Modifier.STATIC);
+				method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER);
+				method.addParameter(this.target, INSTANCE_PARAMETER);
+				method.addCode(generateInjectedElementMethodCode(hints, generatedClass, injectedElement));
+			});
+			ArgumentCodeGenerator argumentCodeGenerator = ArgumentCodeGenerator
+					.of(RegisteredBean.class, REGISTERED_BEAN_PARAMETER).and(this.target, INSTANCE_PARAMETER);
+			return generatedMethod.toMethodReference().toInvokeCodeBlock(argumentCodeGenerator, generatedClass.getName());
+		}
+
+		private CodeBlock generateInjectedElementMethodCode(RuntimeHints hints, GeneratedClass generatedClass,
+				InjectedElement injectedElement) {
+
+			CodeBlock.Builder code = CodeBlock.builder();
+			InjectionCodeGenerator injectionCodeGenerator =
+					new InjectionCodeGenerator(generatedClass.getName(), hints);
+			CodeBlock resourceToInject = generateResourceToInjectCode(generatedClass.getMethods(),
+					(PersistenceElement) injectedElement);
+			code.add(injectionCodeGenerator.generateInjectionCode(
+					injectedElement.getMember(), INSTANCE_PARAMETER,
+					resourceToInject));
+			return code.build();
+		}
+
 		private CodeBlock generateResourceToInjectCode(
 				GeneratedMethods generatedMethods, PersistenceElement injectedElement) {
 
@@ -821,7 +850,7 @@ private CodeBlock generateResourceToInjectCode(
 						EntityManagerFactoryUtils.class, ListableBeanFactory.class,
 						REGISTERED_BEAN_PARAMETER, unitName);
 			}
-			String[] methodNameParts = {"get", unitName, "EntityManager"};
+			String[] methodNameParts = { "get", unitName, "EntityManager" };
 			GeneratedMethod generatedMethod = generatedMethods.add(methodNameParts, method ->
 					generateGetEntityManagerMethod(method, injectedElement));
 			return CodeBlock.of("$L($L)", generatedMethod.getName(), REGISTERED_BEAN_PARAMETER);

spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java

@@ -32,6 +32,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
+import org.springframework.aot.hint.FieldHint;
 import org.springframework.aot.hint.TypeReference;
 import org.springframework.aot.test.generate.TestGenerationContext;
 import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
@@ -160,6 +161,28 @@ void processAheadOfTimeWhenPersistenceContextWithCustomPropertiesOnMethod() {
 		});
 	}
 
+	@Test
+	void processAheadOfTimeWhenPersistenceContextOnPrivateFields() {
+		RegisteredBean registeredBean = registerBean(
+				SeveralPersistenceContextField.class);
+		testCompile(registeredBean, (actual, compiled) -> {
+			EntityManagerFactory entityManagerFactory = mock();
+			this.beanFactory.registerSingleton("custom", entityManagerFactory);
+			this.beanFactory.registerAlias("custom", "another");
+			SeveralPersistenceContextField instance = new SeveralPersistenceContextField();
+			actual.accept(registeredBean, instance);
+			assertThat(instance).extracting("customEntityManager").isNotNull();
+			assertThat(instance).extracting("anotherEntityManager").isNotNull();
+			assertThat(this.generationContext.getRuntimeHints().reflection().typeHints())
+					.singleElement().satisfies(typeHint -> {
+						assertThat(typeHint.getType()).isEqualTo(
+								TypeReference.of(SeveralPersistenceContextField.class));
+						assertThat(typeHint.fields().map(FieldHint::getName))
+								.containsOnly("customEntityManager", "anotherEntityManager");
+					});
+		});
+	}
+
 	private RegisteredBean registerBean(Class<?> beanClass) {
 		String beanName = "testBean";
 		this.beanFactory.registerBeanDefinition(beanName,
@@ -256,4 +279,16 @@ public void setEntityManager(EntityManager entityManager) {
 
 	}
 
+	static class SeveralPersistenceContextField {
+
+		@SuppressWarnings("unused")
+		@PersistenceContext(name = "custom")
+		private EntityManager customEntityManager;
+
+		@SuppressWarnings("unused")
+		@PersistenceContext(name = "another")
+		private EntityManager anotherEntityManager;
+
+	}
+
 }