Add support for JPA 4.0 @PersistenceAgent injection

Closes gh-36264

spring-orm/src/main/java/org/springframework/orm/jpa/AbstractEntityManagerFactoryBean.java

@@ -261,6 +261,11 @@ public abstract class AbstractEntityManagerFactoryBean implements
 		return this.entityManagerInterface;
 	}
 
+	@Override
+	public @Nullable Class<?> getEntityAgentInterface() {
+		return this.entityAgentInterface;
+	}
+
 	/**
 	 * Specify the vendor-specific JpaDialect implementation to associate with
 	 * this EntityManagerFactory. This will be exposed through the
@@ -419,8 +424,7 @@ public abstract class AbstractEntityManagerFactoryBean implements
 
 		this.sharedEntityManager = SharedEntityManagerCreator.createSharedEntityManager(this.entityManagerFactory);
 		if (this.entityAgentInterface != null) {
-			this.sharedEntityAgent = SharedEntityManagerCreator.createSharedEntityAgent(
-					this.entityManagerFactory, null, this.entityAgentInterface);
+			this.sharedEntityAgent = SharedEntityManagerCreator.createSharedEntityAgent(this.entityManagerFactory);
 		}
 	}
 

spring-orm/src/main/java/org/springframework/orm/jpa/EntityManagerFactoryInfo.java

@@ -83,6 +83,16 @@ public interface EntityManagerFactoryInfo {
 	 */
 	@Nullable Class<? extends EntityManager> getEntityManagerInterface();
 
+	/**
+	 * Return the (potentially vendor-specific) EntityAgent interface
+	 * that this factory's EntityAgents will implement.
+	 * <p>A {@code null} return value suggests that autodetection is supposed
+	 * to happen: either based on a target {@code EntityAgent} instance
+	 * or simply defaulting to {@code jakarta.persistence.EntityAgent}.
+	 * @since 7.1
+	 */
+	@Nullable Class<?> getEntityAgentInterface();
+
 	/**
 	 * Return the vendor-specific JpaDialect implementation for this
 	 * EntityManagerFactory, or {@code null} if not known.

spring-orm/src/main/java/org/springframework/orm/jpa/EntityManagerFactoryUtils.java

@@ -380,10 +380,12 @@ public abstract class EntityManagerFactoryUtils {
 
 	/**
 	 * Determine the {@code jakarta.persistence.EntityAgent} class.
-	 * @return the {@code EntityAgent} class, or {@code null} if not available
-	 * @since 7.0.4
+	 * <p>Intended for internal adaptation to JPA 3.2 vs 4.0.
+	 * @return the {@code EntityAgent} class (on JPA 4.0),
+	 * or {@code null} if not available (on JPA 3.2)
+	 * @since 7.1
 	 */
-	static @Nullable Class<?> getEntityAgentClass() {
+	public static @Nullable Class<?> getEntityAgentClass() {
 		return (CREATE_ENTITY_AGENT_METHOD != null ? CREATE_ENTITY_AGENT_METHOD.getReturnType() : null);
 	}
 

spring-orm/src/main/java/org/springframework/orm/jpa/SharedEntityManagerCreator.java

@@ -38,6 +38,7 @@ import org.apache.commons.logging.LogFactory;
 import org.jspecify.annotations.Nullable;
 
 import org.springframework.transaction.support.TransactionSynchronizationManager;
+import org.springframework.util.Assert;
 import org.springframework.util.ClassUtils;
 import org.springframework.util.CollectionUtils;
 import org.springframework.util.ConcurrentReferenceHashMap;
@@ -177,19 +178,38 @@ public abstract class SharedEntityManagerCreator {
 				ifcs, new SharedEntityManagerInvocationHandler(emf, properties, synchronizedWithTransaction));
 	}
 
+	/**
+	 * Create a transactional EntityAgent proxy for the given EntityManagerFactory.
+	 * @param emf the EntityManagerFactory to obtain EntityAgents from as needed
+	 * @return a shareable transactional EntityAgent proxy
+	 * (typed to {@code Object} for compatibility with JPA 3.2; this will change
+	 * to {@code EntityAgent} once Spring establishes a JPA 4.0 minimum baseline)
+	 * @since 7.1
+	 */
+	public static Object createSharedEntityAgent(EntityManagerFactory emf) {
+		return createSharedEntityAgent(emf, null);
+	}
+
 	/**
 	 * Create a transactional EntityAgent proxy for the given EntityManagerFactory.
 	 * @param emf the EntityManagerFactory to obtain EntityAgentss from as needed
 	 * @param properties the properties to be passed into the
 	 * {@code createEntityAgent} call (may be {@code null})
-	 * @param ifc the interfaces to be implemented by the EntityAgent
 	 * @return a shareable transactional EntityAgent proxy
-	 * @since 7.0.4
+	 * (typed to {@code Object} for compatibility with JPA 3.2; this will change
+	 * to {@code EntityAgent} once Spring establishes a JPA 4.0 minimum baseline)
+	 * @since 7.1
 	 */
-	static Object createSharedEntityAgent(EntityManagerFactory emf, @Nullable Map<?, ?> properties, Class<?> ifc) {
+	public static Object createSharedEntityAgent(EntityManagerFactory emf, @Nullable Map<?, ?> properties) {
 		ClassLoader cl = null;
+		Class<?> ifc = null;
 		if (emf instanceof EntityManagerFactoryInfo emfInfo) {
 			cl = emfInfo.getBeanClassLoader();
+			ifc = emfInfo.getEntityAgentInterface();
+		}
+		if (ifc == null) {
+			ifc = EntityManagerFactoryUtils.getEntityAgentClass();
+			Assert.state(ifc != null, "JPA 4.0 EntityAgent class not found");
 		}
 		return Proxy.newProxyInstance(
 				(cl != null ? cl : SharedEntityManagerCreator.class.getClassLoader()),

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

@@ -18,11 +18,11 @@ package org.springframework.orm.jpa.support;
 
 import java.beans.PropertyDescriptor;
 import java.io.Serializable;
+import java.lang.annotation.Annotation;
 import java.lang.reflect.AnnotatedElement;
 import java.lang.reflect.Member;
 import java.lang.reflect.Modifier;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -190,6 +190,13 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 		DestructionAwareBeanPostProcessor, MergedBeanDefinitionPostProcessor, BeanRegistrationAotProcessor,
 		PriorityOrdered, BeanFactoryAware, Serializable {
 
+	private static final @Nullable Class<? extends Annotation> PERSISTENCE_AGENT_ANNOTATION;
+
+	static {
+		PERSISTENCE_AGENT_ANNOTATION = AnnotationUtils.loadAnnotationType("jakarta.persistence.PersistenceAgent");
+	}
+
+
 	private @Nullable Object jndiEnvironment;
 
 	private boolean resourceRef = true;
@@ -200,6 +207,8 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 
 	private transient @Nullable Map<String, String> extendedPersistenceContexts;
 
+	private transient @Nullable Map<String, String> persistenceAgents;
+
 	private transient String defaultPersistenceUnitName = "";
 
 	private int order = Ordered.LOWEST_PRECEDENCE - 4;
@@ -305,6 +314,20 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 		this.extendedPersistenceContexts = extendedPersistenceContexts;
 	}
 
+	/**
+	 * Specify the persistence agent locations for EntityAgent lookups,
+	 * as a Map from persistence unit name to persistence agent JNDI name
+	 * (which needs to resolve to an EntityAgent instance).
+	 * <p>In case of no unit name specified in the annotation, the specified value
+	 * for the {@link #setDefaultPersistenceUnitName default persistence unit}
+	 * will be taken (by default, the value mapped to the empty String),
+	 * or simply the single persistence unit if there is only one.
+	 * @since 7.1
+	 */
+	public void setPersistenceAgents(@Nullable Map<String, String> persistenceAgents) {
+		this.persistenceAgents = persistenceAgents;
+	}
+
 	/**
 	 * Specify the default persistence unit name, to be used in case
 	 * of no unit name specified in an {@code @PersistenceUnit} /
@@ -417,7 +440,7 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 	}
 
 	private InjectionMetadata buildPersistenceMetadata(Class<?> clazz) {
-		if (!AnnotationUtils.isCandidateClass(clazz, Arrays.asList(PersistenceContext.class, PersistenceUnit.class))) {
+		if (!AnnotationUtils.isCandidateClass(clazz, PersistenceUnit.class)) {
 			return InjectionMetadata.EMPTY;
 		}
 
@@ -428,8 +451,9 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 			final List<InjectionMetadata.InjectedElement> currElements = new ArrayList<>();
 
 			ReflectionUtils.doWithLocalFields(targetClass, field -> {
-				if (field.isAnnotationPresent(PersistenceContext.class) ||
-						field.isAnnotationPresent(PersistenceUnit.class)) {
+				if (field.isAnnotationPresent(PersistenceUnit.class) ||
+						field.isAnnotationPresent(PersistenceContext.class) ||
+						(PERSISTENCE_AGENT_ANNOTATION != null && field.isAnnotationPresent(PERSISTENCE_AGENT_ANNOTATION))) {
 					if (Modifier.isStatic(field.getModifiers())) {
 						throw new IllegalStateException("Persistence annotations are not supported on static fields");
 					}
@@ -441,8 +465,9 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 				if (method.isBridge()) {
 					return;
 				}
-				if ((method.isAnnotationPresent(PersistenceContext.class) ||
-						method.isAnnotationPresent(PersistenceUnit.class)) &&
+				if ((method.isAnnotationPresent(PersistenceUnit.class) ||
+						method.isAnnotationPresent(PersistenceContext.class) ||
+						(PERSISTENCE_AGENT_ANNOTATION != null && method.isAnnotationPresent(PERSISTENCE_AGENT_ANNOTATION))) &&
 						method.equals(BridgeMethodResolver.getMostSpecificMethod(method, clazz))) {
 					if (Modifier.isStatic(method.getModifiers())) {
 						throw new IllegalStateException("Persistence annotations are not supported on static methods");
@@ -525,6 +550,37 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 		return null;
 	}
 
+	/**
+	 * Return a specified persistence context for the given unit name, as defined
+	 * through the "persistenceAgents" map.
+	 * @param unitName the name of the persistence unit
+	 * @return the corresponding EntityAgent, or {@code null} if none found
+	 * @since 7.1
+	 * @see #setPersistenceAgents
+	 */
+	private @Nullable Object getPersistenceAgent(@Nullable String unitName) {
+		Class<?> entityAgentClass = EntityManagerFactoryUtils.getEntityAgentClass();
+		if (entityAgentClass != null && this.persistenceAgents != null) {
+			String unitNameForLookup = (unitName != null ? unitName : "");
+			if (unitNameForLookup.isEmpty()) {
+				unitNameForLookup = this.defaultPersistenceUnitName;
+			}
+			String jndiName = this.persistenceAgents.get(unitNameForLookup);
+			if (jndiName == null && unitNameForLookup.isEmpty() && this.persistenceAgents.size() == 1) {
+				jndiName = this.persistenceAgents.values().iterator().next();
+			}
+			if (jndiName != null) {
+				try {
+					return lookup(jndiName, entityAgentClass);
+				}
+				catch (Exception ex) {
+					throw new IllegalStateException("Could not obtain EntityAgent [" + jndiName + "] from JNDI", ex);
+				}
+			}
+		}
+		return null;
+	}
+
 	/**
 	 * Find an EntityManagerFactory with the given name in the current Spring
 	 * application context, falling back to a single default EntityManagerFactory
@@ -636,6 +692,8 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 	 */
 	private class PersistenceElement extends InjectionMetadata.InjectedElement {
 
+		private final boolean persistenceUnit;
+
 		private final String unitName;
 
 		private @Nullable PersistenceContextType type;
@@ -646,48 +704,69 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 
 		public PersistenceElement(Member member, AnnotatedElement ae, @Nullable PropertyDescriptor pd) {
 			super(member, pd);
-			PersistenceContext pc = ae.getAnnotation(PersistenceContext.class);
 			PersistenceUnit pu = ae.getAnnotation(PersistenceUnit.class);
-			Class<?> resourceType = EntityManager.class;
-			if (pc != null) {
+			PersistenceContext pc = ae.getAnnotation(PersistenceContext.class);
+			Class<?> entityAgentClass = EntityManagerFactoryUtils.getEntityAgentClass();
+			Annotation pa = (PERSISTENCE_AGENT_ANNOTATION != null ?
+					ae.getAnnotation(PERSISTENCE_AGENT_ANNOTATION) : null);
+			Class<?> resourceType;
+			if (entityAgentClass != null && pa != null) {
+				if (pu != null) {
+					throw new IllegalStateException("Member may only be annotated with either " +
+							"@PersistenceAgent or @PersistenceUnit, not both: " + member);
+				}
+				resourceType = entityAgentClass;
+				this.persistenceUnit = false;
+				this.unitName = String.valueOf(AnnotationUtils.getValue(pa, "unitName"));
+				this.properties = parseProperties((PersistenceProperty[]) AnnotationUtils.getValue(pa, "properties"));
+			}
+			else if (pc != null) {
 				if (pu != null) {
 					throw new IllegalStateException("Member may only be annotated with either " +
 							"@PersistenceContext or @PersistenceUnit, not both: " + member);
 				}
-				Properties properties = null;
-				PersistenceProperty[] pps = pc.properties();
-				if (!ObjectUtils.isEmpty(pps)) {
-					properties = new Properties();
-					for (PersistenceProperty pp : pps) {
-						properties.setProperty(pp.name(), pp.value());
-					}
-				}
+				resourceType = EntityManager.class;
+				this.persistenceUnit = false;
 				this.unitName = pc.unitName();
 				this.type = pc.type();
 				this.synchronizedWithTransaction = SynchronizationType.SYNCHRONIZED.equals(pc.synchronization());
-				this.properties = properties;
+				this.properties = parseProperties(pc.properties());
 			}
 			else {
 				resourceType = EntityManagerFactory.class;
+				this.persistenceUnit = true;
 				this.unitName = pu.unitName();
 			}
 			checkResourceType(resourceType);
 		}
 
+		private static @Nullable Properties parseProperties(PersistenceProperty @Nullable [] pps) {
+			if (ObjectUtils.isEmpty(pps)) {
+				return null;
+			}
+			Properties properties = new Properties();
+			for (PersistenceProperty pp : pps) {
+				properties.setProperty(pp.name(), pp.value());
+			}
+			return properties;
+		}
+
 		/**
 		 * Resolve the object against the application context.
 		 */
 		@Override
 		protected Object getResourceToInject(Object target, @Nullable String requestingBeanName) {
 			// Resolves to EntityManagerFactory or EntityManager.
-			if (this.type != null) {
+			if (this.persistenceUnit) {
+				return resolveEntityManagerFactory(requestingBeanName);
+			}
+			else if (this.type != null) {
 				return (this.type == PersistenceContextType.EXTENDED ?
 						resolveExtendedEntityManager(target, requestingBeanName) :
 						resolveEntityManager(requestingBeanName));
 			}
 			else {
-				// OK, so we need an EntityManagerFactory...
-				return resolveEntityManagerFactory(requestingBeanName);
+				return resolveEntityAgent(requestingBeanName);
 			}
 		}
 
@@ -749,6 +828,23 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 			}
 			return em;
 		}
+
+		private Object resolveEntityAgent(@Nullable String requestingBeanName) {
+			// Obtain EntityAgent reference from JNDI?
+			Object ea = getPersistenceAgent(this.unitName);
+			if (ea == null) {
+				// No pre-built EntityManager found -> build one based on factory.
+				// Obtain EntityManagerFactory from JNDI?
+				EntityManagerFactory emf = getPersistenceUnit(this.unitName);
+				if (emf == null) {
+					// Need to search for EntityManagerFactory beans.
+					emf = findEntityManagerFactory(this.unitName, requestingBeanName);
+				}
+				// Inject a shared transactional EntityManager proxy.
+				ea = SharedEntityManagerCreator.createSharedEntityAgent(emf, this.properties);
+			}
+			return ea;
+		}
 	}
 
 
@@ -833,25 +929,35 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 				GeneratedMethods generatedMethods, PersistenceElement injectedElement) {
 
 			String unitName = injectedElement.unitName;
-			boolean requireEntityManager = (injectedElement.type != null);
-			if (!requireEntityManager) {
+			if (injectedElement.persistenceUnit) {
 				return CodeBlock.of(
 						"$T.findEntityManagerFactory(($T) $L.getBeanFactory(), $S)",
 						EntityManagerFactoryUtils.class, ListableBeanFactory.class,
 						REGISTERED_BEAN_PARAMETER, unitName);
 			}
-			String[] methodNameParts = { "get", unitName, "EntityManager" };
-			GeneratedMethod generatedMethod = generatedMethods.add(methodNameParts, method ->
-					generateGetEntityManagerMethod(method, injectedElement));
-			return CodeBlock.of("$L($L)", generatedMethod.getName(), REGISTERED_BEAN_PARAMETER);
+			else if (injectedElement.type != null) {
+				if (injectedElement.type == PersistenceContextType.EXTENDED) {
+					throw new UnsupportedOperationException(
+							"PersistenceContextType.EXTENDED not supported for AOT processing");
+				}
+				String[] methodNameParts = {"get", unitName, "EntityManager"};
+				GeneratedMethod generatedMethod = generatedMethods.add(methodNameParts, method ->
+						generateGetEntityManagerMethod(method, injectedElement));
+				return CodeBlock.of("$L($L)", generatedMethod.getName(), REGISTERED_BEAN_PARAMETER);
+			}
+			else {
+				String[] methodNameParts = {"get", unitName, "EntityAgent"};
+				GeneratedMethod generatedMethod = generatedMethods.add(methodNameParts, method ->
+						generateGetEntityAgentMethod(method, injectedElement));
+				return CodeBlock.of("$L($L)", generatedMethod.getName(), REGISTERED_BEAN_PARAMETER);
+			}
 		}
 
-		@SuppressWarnings("NullAway") // Dataflow analysis limitation
 		private void generateGetEntityManagerMethod(MethodSpec.Builder method, PersistenceElement injectedElement) {
 			String unitName = injectedElement.unitName;
 			Properties properties = injectedElement.properties;
 			method.addJavadoc("Get the '$L' {@link $T}.",
-					(StringUtils.hasLength(unitName)) ? unitName : "default",
+					(StringUtils.hasLength(unitName) ? unitName : "default"),
 					EntityManager.class);
 			method.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
 					javax.lang.model.element.Modifier.STATIC);
@@ -861,10 +967,8 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 					"$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) {
-				method.addStatement("$T properties = new Properties()",
-						Properties.class);
+			if (properties != null) {
+				method.addStatement("$T properties = new Properties()", Properties.class);
 				for (String propertyName : new TreeSet<>(properties.stringPropertyNames())) {
 					method.addStatement("properties.put($S, $S)", propertyName, properties.getProperty(propertyName));
 				}
@@ -872,9 +976,38 @@ public class PersistenceAnnotationBeanPostProcessor implements InstantiationAwar
 			method.addStatement(
 					"return $T.createSharedEntityManager(entityManagerFactory, $L, $L)",
 					SharedEntityManagerCreator.class,
-					(hasProperties) ? "properties" : null,
+					(properties != null ? "properties" : null),
 					injectedElement.synchronizedWithTransaction);
 		}
+
+		private void generateGetEntityAgentMethod(MethodSpec.Builder method, PersistenceElement injectedElement) {
+			Class<?> entityAgentClass = EntityManagerFactoryUtils.getEntityAgentClass();
+			Assert.state(entityAgentClass != null, "JPA 4.0 EntityAgent class not found");
+			String unitName = injectedElement.unitName;
+			Properties properties = injectedElement.properties;
+			method.addJavadoc("Get the '$L' {@link $T}.",
+					(StringUtils.hasLength(unitName) ? unitName : "default"),
+					entityAgentClass);
+			method.addModifiers(javax.lang.model.element.Modifier.PUBLIC,
+					javax.lang.model.element.Modifier.STATIC);
+			method.returns(entityAgentClass);
+			method.addParameter(RegisteredBean.class, REGISTERED_BEAN_PARAMETER);
+			method.addStatement(
+					"$T entityManagerFactory = $T.findEntityManagerFactory(($T) $L.getBeanFactory(), $S)",
+					EntityManagerFactory.class, EntityManagerFactoryUtils.class,
+					ListableBeanFactory.class, REGISTERED_BEAN_PARAMETER, unitName);
+			if (properties != null) {
+				method.addStatement("$T properties = new Properties()",
+						Properties.class);
+				for (String propertyName : new TreeSet<>(properties.stringPropertyNames())) {
+					method.addStatement("properties.put($S, $S)", propertyName, properties.getProperty(propertyName));
+				}
+			}
+			method.addStatement(
+					"return $T.createSharedEntityAgent(entityManagerFactory, $L)",
+					SharedEntityManagerCreator.class,
+					(properties != null ? "properties" : null));
+		}
 	}
 
 }