Update PersistenceAnnotationBeanPostProcessor AOT support

Update `PersistenceAnnotationBeanPostProcessor` so that it provides
AOT contributions via the `BeanRegistrationAotProcessor` interface.

See gh-28414
This commit is contained in:
Phillip Webb
2022-04-26 20:48:40 -07:00
parent e4a8258fb2
commit b677eb90f9
4 changed files with 690 additions and 1 deletions

View File

@@ -0,0 +1,116 @@
/*
* Copyright 2002-2022 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.orm.jpa.support;
import java.lang.reflect.Field;
import java.lang.reflect.Member;
import java.lang.reflect.Method;
import org.springframework.aot.generate.AccessVisibility;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.javapoet.CodeBlock;
import org.springframework.util.Assert;
import org.springframework.util.ReflectionUtils;
/**
* Internal code generator that can inject a value into a field or single-arg
* method.
* <p>
* Generates code in the form:<pre class="code">{@code
* instance.age = value;
* }</pre> or <pre class="code">{@code
* instance.setAge(value);
* }</pre>
* <p>
* Will also generate reflection based injection and register hints if the
* member is not visible.
*
* @author Phillip Webb
* @since 6.0
*/
class InjectionCodeGenerator {
private final RuntimeHints hints;
InjectionCodeGenerator(RuntimeHints hints) {
Assert.notNull(hints, "Hints must not be null");
this.hints = hints;
}
CodeBlock generateInjectionCode(Member member, String instanceVariable,
CodeBlock resourceToInject) {
if (member instanceof Field field) {
return generateFieldInjectionCode(field, instanceVariable, resourceToInject);
}
if (member instanceof Method method) {
return generateMethodInjectionCode(method, instanceVariable,
resourceToInject);
}
throw new IllegalStateException(
"Unsupported member type " + member.getClass().getName());
}
private CodeBlock generateFieldInjectionCode(Field field, String instanceVariable,
CodeBlock resourceToInject) {
CodeBlock.Builder builder = CodeBlock.builder();
AccessVisibility visibility = AccessVisibility.forMember(field);
if (visibility == AccessVisibility.PRIVATE
|| visibility == AccessVisibility.PROTECTED) {
this.hints.reflection().registerField(field);
builder.addStatement("$T field = $T.findField($T.class, $S)", Field.class,
ReflectionUtils.class, field.getDeclaringClass(), field.getName());
builder.addStatement("$T.makeAccessible($L)", ReflectionUtils.class, "field");
builder.addStatement("$T.setField($L, $L, $L)", ReflectionUtils.class,
"field", instanceVariable, resourceToInject);
}
else {
builder.addStatement("$L.$L = $L", instanceVariable, field.getName(),
resourceToInject);
}
return builder.build();
}
private CodeBlock generateMethodInjectionCode(Method method, String instanceVariable,
CodeBlock resourceToInject) {
Assert.isTrue(method.getParameterCount() == 1,
"Method '" + method.getName() + "' must declare a single parameter");
CodeBlock.Builder builder = CodeBlock.builder();
AccessVisibility visibility = AccessVisibility.forMember(method);
if (visibility == AccessVisibility.PRIVATE
|| visibility == AccessVisibility.PROTECTED) {
this.hints.reflection().registerMethod(method);
builder.addStatement("$T method = $T.findMethod($T.class, $S, $T.class)",
Method.class, ReflectionUtils.class, method.getDeclaringClass(),
method.getName(), method.getParameterTypes()[0]);
builder.addStatement("$T.makeAccessible($L)", ReflectionUtils.class,
"method");
builder.addStatement("$T.invokeMethod($L, $L, $L)", ReflectionUtils.class,
"method", instanceVariable, resourceToInject);
}
else {
builder.addStatement("$L.$L($L)", instanceVariable, method.getName(),
resourceToInject);
}
return builder.build();
}
}

View File

@@ -29,6 +29,7 @@ import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
import jakarta.persistence.EntityManager;
@@ -39,8 +40,15 @@ import jakarta.persistence.PersistenceProperty;
import jakarta.persistence.PersistenceUnit;
import jakarta.persistence.SynchronizationType;
import org.springframework.aot.generate.GeneratedMethod;
import org.springframework.aot.generate.GeneratedMethods;
import org.springframework.aot.generate.GenerationContext;
import org.springframework.aot.generate.MethodGenerator;
import org.springframework.aot.generate.MethodNameGenerator;
import org.springframework.aot.generate.MethodReference;
import org.springframework.aot.generator.CodeContribution;
import org.springframework.aot.generator.ProtectedAccess.Options;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.PropertyValues;
import org.springframework.beans.factory.BeanCreationException;
@@ -50,6 +58,9 @@ import org.springframework.beans.factory.ListableBeanFactory;
import org.springframework.beans.factory.NoSuchBeanDefinitionException;
import org.springframework.beans.factory.annotation.InjectionMetadata;
import org.springframework.beans.factory.annotation.InjectionMetadata.InjectedElement;
import org.springframework.beans.factory.aot.BeanRegistrationAotContribution;
import org.springframework.beans.factory.aot.BeanRegistrationAotProcessor;
import org.springframework.beans.factory.aot.BeanRegistrationCode;
import org.springframework.beans.factory.config.ConfigurableBeanFactory;
import org.springframework.beans.factory.config.ConfigurableListableBeanFactory;
import org.springframework.beans.factory.config.DestructionAwareBeanPostProcessor;
@@ -59,12 +70,17 @@ import org.springframework.beans.factory.generator.AotContributingBeanPostProces
import org.springframework.beans.factory.generator.BeanFieldGenerator;
import org.springframework.beans.factory.generator.BeanInstantiationContribution;
import org.springframework.beans.factory.support.MergedBeanDefinitionPostProcessor;
import org.springframework.beans.factory.support.RegisteredBean;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.core.BridgeMethodResolver;
import org.springframework.core.Ordered;
import org.springframework.core.PriorityOrdered;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.javapoet.ClassName;
import org.springframework.javapoet.CodeBlock;
import org.springframework.javapoet.JavaFile;
import org.springframework.javapoet.MethodSpec;
import org.springframework.javapoet.TypeSpec;
import org.springframework.javapoet.support.MultiStatement;
import org.springframework.jndi.JndiLocatorDelegate;
import org.springframework.jndi.JndiTemplate;
@@ -176,6 +192,7 @@ import org.springframework.util.StringUtils;
* @author Rod Johnson
* @author Juergen Hoeller
* @author Stephane Nicoll
* @author Phillip Webb
* @since 2.0
* @see jakarta.persistence.PersistenceUnit
* @see jakarta.persistence.PersistenceContext
@@ -183,7 +200,7 @@ import org.springframework.util.StringUtils;
@SuppressWarnings("serial")
public class PersistenceAnnotationBeanPostProcessor
implements InstantiationAwareBeanPostProcessor, DestructionAwareBeanPostProcessor,
MergedBeanDefinitionPostProcessor, AotContributingBeanPostProcessor,
MergedBeanDefinitionPostProcessor, AotContributingBeanPostProcessor, BeanRegistrationAotProcessor,
PriorityOrdered, BeanFactoryAware, Serializable {
@Nullable
@@ -358,6 +375,19 @@ public class PersistenceAnnotationBeanPostProcessor
return null;
}
@Override
public BeanRegistrationAotContribution processAheadOfTime(RegisteredBean registeredBean) {
Class<?> beanClass = registeredBean.getBeanClass();
String beanName = registeredBean.getBeanName();
RootBeanDefinition beanDefinition = registeredBean.getMergedBeanDefinition();
InjectionMetadata metadata = findInjectionMetadata(beanDefinition, beanClass, beanName);
Collection<InjectedElement> injectedElements = metadata.getInjectedElements();
if (!CollectionUtils.isEmpty(injectedElements)) {
return new AotContribution(beanClass, injectedElements);
}
return null;
}
private InjectionMetadata findInjectionMetadata(RootBeanDefinition beanDefinition, Class<?> beanType, String beanName) {
InjectionMetadata metadata = findPersistenceMetadata(beanName, beanType, null);
metadata.checkConfigMembers(beanDefinition);
@@ -815,4 +845,126 @@ public class PersistenceAnnotationBeanPostProcessor
}
private static class AotContribution implements BeanRegistrationAotContribution {
private static final String APPLY_METHOD = "apply";
private static final String REGISTERED_BEAN_PARAMETER = "registeredBean";
private static final String INSTANCE_PARAMETER = "instance";
private final Class<?> target;
private final Collection<InjectedElement> injectedElements;
AotContribution(Class<?> target, Collection<InjectedElement> injectedElements) {
this.target = target;
this.injectedElements = injectedElements;
}
@Override
public void applyTo(GenerationContext generationContext,
BeanRegistrationCode beanRegistrationCode) {
ClassName className = generationContext.getClassNameGenerator()
.generateClassName(this.target, "PersistenceInjection");
TypeSpec.Builder classBuilder = TypeSpec.classBuilder(className);
classBuilder.addJavadoc("Persistence injection for {@link $T}.", this.target);
classBuilder.addModifiers(javax.lang.model.element.Modifier.PUBLIC);
GeneratedMethods methods = new GeneratedMethods(
new MethodNameGenerator(APPLY_METHOD));
classBuilder.addMethod(generateMethod(generationContext.getRuntimeHints(),
className, methods));
methods.doWithMethodSpecs(classBuilder::addMethod);
JavaFile javaFile = JavaFile
.builder(className.packageName(), classBuilder.build()).build();
generationContext.getGeneratedFiles().addSourceFile(javaFile);
beanRegistrationCode.addInstancePostProcessor(
MethodReference.ofStatic(className, APPLY_METHOD));
}
private MethodSpec generateMethod(RuntimeHints hints, ClassName className,
MethodGenerator methodGenerator) {
MethodSpec.Builder builder = MethodSpec.methodBuilder(APPLY_METHOD);
builder.addJavadoc("Apply the persistence injection.");
builder.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
javax.lang.model.element.Modifier.STATIC);
builder.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER);
builder.addParameter(this.target, INSTANCE_PARAMETER);
builder.returns(this.target);
builder.addCode(generateMethodCode(hints, methodGenerator));
return builder.build();
}
private CodeBlock generateMethodCode(RuntimeHints hints,
MethodGenerator methodGenerator) {
CodeBlock.Builder builder = CodeBlock.builder();
InjectionCodeGenerator injectionCodeGenerator = new InjectionCodeGenerator(
hints);
for (InjectedElement injectedElement : this.injectedElements) {
CodeBlock resourceToInject = getResourceToInject(methodGenerator,
(PersistenceElement) injectedElement);
builder.add(injectionCodeGenerator.generateInjectionCode(
injectedElement.getMember(), INSTANCE_PARAMETER,
resourceToInject));
}
builder.addStatement("return $L", INSTANCE_PARAMETER);
return builder.build();
}
private CodeBlock getResourceToInject(MethodGenerator methodGenerator,
PersistenceElement injectedElement) {
String unitName = injectedElement.unitName;
boolean requireEntityManager = (injectedElement.type != null);
if (!requireEntityManager) {
return CodeBlock.of(
"$T.findEntityManagerFactory(($T) $L.getBeanFactory(), $S)",
EntityManagerFactoryUtils.class, ListableBeanFactory.class,
REGISTERED_BEAN_PARAMETER, unitName);
}
GeneratedMethod getEntityManagerMethod = methodGenerator
.generateMethod("get", unitName, "EntityManager")
.using(builder -> buildGetEntityManagerMethod(builder,
injectedElement));
return CodeBlock.of("$L($L)", getEntityManagerMethod.getName(),
REGISTERED_BEAN_PARAMETER);
}
private void buildGetEntityManagerMethod(MethodSpec.Builder builder,
PersistenceElement injectedElement) {
String unitName = injectedElement.unitName;
Properties properties = injectedElement.properties;
builder.addJavadoc("Get the '$L' {@link $T}",
(StringUtils.hasLength(unitName)) ? unitName : "default",
EntityManager.class);
builder.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
javax.lang.model.element.Modifier.STATIC);
builder.returns(EntityManager.class);
builder.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER);
builder.addStatement(
"$T entityManagerFactory = $T.findEntityManagerFactory(($T) $L.getBeanFactory(), $S)",
EntityManagerFactory.class, EntityManagerFactoryUtils.class,
ListableBeanFactory.class, REGISTERED_BEAN_PARAMETER, unitName);
boolean hasProperties = !CollectionUtils.isEmpty(properties);
if (hasProperties) {
builder.addStatement("$T properties = new Properties()",
Properties.class);
for (String propertyName : new TreeSet<>(
properties.stringPropertyNames())) {
builder.addStatement("properties.put($S, $S)", propertyName,
properties.getProperty(propertyName));
}
}
builder.addStatement(
"return $T.createSharedEntityManager(entityManagerFactory, $L, $L)",
SharedEntityManagerCreator.class,
(hasProperties) ? "properties" : null,
injectedElement.synchronizedWithTransaction);
}
}
}