From 120c228b7064d606163fae4067602b7e8d868ae1 Mon Sep 17 00:00:00 2001 From: Stephane Nicoll Date: Wed, 10 May 2023 13:43:26 +0200 Subject: [PATCH] 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 --- ...ersistenceAnnotationBeanPostProcessor.java | 51 +++++++++++++++---- ...BeanPostProcessorAotContributionTests.java | 35 +++++++++++++ 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/spring-orm/src/main/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessor.java b/spring-orm/src/main/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessor.java index 0c7735cef0c..944b5513c7e 100644 --- a/spring-orm/src/main/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessor.java +++ b/spring-orm/src/main/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessor.java @@ -43,6 +43,7 @@ import org.springframework.aot.generate.GeneratedClass; 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 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar private final Class target; - private final Collection injectedElements; + private final List injectedElements; AotContribution(Class target, Collection injectedElements) { this.target = target; - this.injectedElements = injectedElements; + this.injectedElements = List.copyOf(injectedElements); } @Override @@ -797,19 +798,47 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar 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 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar 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); diff --git a/spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java b/spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java index 4b779f972d7..8e0dbfe12ab 100644 --- a/spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java +++ b/spring-orm/src/test/java/org/springframework/orm/jpa/support/PersistenceAnnotationBeanPostProcessorAotContributionTests.java @@ -32,6 +32,7 @@ import org.assertj.core.api.InstanceOfAssertFactories; 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 @@ class PersistenceAnnotationBeanPostProcessorAotContributionTests { }); } + @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 @@ class PersistenceAnnotationBeanPostProcessorAotContributionTests { } + static class SeveralPersistenceContextField { + + @SuppressWarnings("unused") + @PersistenceContext(name = "custom") + private EntityManager customEntityManager; + + @SuppressWarnings("unused") + @PersistenceContext(name = "another") + private EntityManager anotherEntityManager; + + } + }