Spring Data JPA doesn't boot when using Hibernate with multi-tenancy

Issue: #3425
When Hibernate is configured with multi-tenancy, upon startup
Spring JPA calls PersistenceProvider.fromEntityManager(entityManager)
which initalizes Hibernate Session. This requires a tenant to be present
and may produce failures.

Signed-off-by: Ariel Morelli Andres <[email protected]>

Author: ari-spf

spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/JpaClassUtils.java

@@ -15,7 +15,7 @@
  */
 package org.springframework.data.jpa.provider;
 
-import jakarta.persistence.EntityManager;
+import jakarta.persistence.EntityManagerFactory;
 import jakarta.persistence.metamodel.Metamodel;
 
 import org.springframework.lang.Nullable;
@@ -30,6 +30,7 @@
  * @author Jens Schauder
  * @author Donghun Shin
  * @author Greg Turnquist
+ * @author Ariel Morelli Andres, Atlassian US, Inc
  */
 abstract class JpaClassUtils {
 
@@ -39,19 +40,15 @@ abstract class JpaClassUtils {
 	private JpaClassUtils() {}
 
 	/**
-	 * Returns whether the given {@link EntityManager} is of the given type.
+	 * Returns whether the given {@link EntityManagerFactory} is of the given type.
 	 *
-	 * @param em must not be {@literal null}.
-	 * @param type the fully qualified expected {@link EntityManager} type, must not be {@literal null} or empty.
-	 * @return whether the given {@code EntityManager} is of the given type.
+	 * @param factory must not be {@literal null}.
+	 * @param type the fully qualified expected {@link EntityManagerFactory} type, must not be {@literal null} or empty.
+	 * @return whether the given {@code EntityManagerFactory} is of the given type.
 	 */
-	public static boolean isEntityManagerOfType(EntityManager em, String type) {
+	public static boolean isEntityManagerFactoryOfType(EntityManagerFactory factory, String type) {
 
-		EntityManager entityManagerToUse = em.getDelegate() instanceof EntityManager delegate //
-				? delegate //
-				: em;
-
-		return isOfType(entityManagerToUse, type, entityManagerToUse.getClass().getClassLoader());
+		return isOfType(factory, type, factory.getClass().getClassLoader());
 	}
 
 	public static boolean isMetamodelOfType(Metamodel metamodel, String type) {

spring-data-jpa/src/main/java/org/springframework/data/jpa/provider/PersistenceProvider.java

@@ -19,15 +19,18 @@
 import static org.springframework.data.jpa.provider.PersistenceProvider.Constants.*;
 
 import jakarta.persistence.EntityManager;
+import jakarta.persistence.EntityManagerFactory;
 import jakarta.persistence.Query;
 import jakarta.persistence.metamodel.IdentifiableType;
 import jakarta.persistence.metamodel.Metamodel;
 import jakarta.persistence.metamodel.SingularAttribute;
 
+import java.lang.reflect.Proxy;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.NoSuchElementException;
+import java.util.Objects;
 import java.util.Set;
 
 import org.eclipse.persistence.config.QueryHints;
@@ -52,6 +55,7 @@
  * @author Jens Schauder
  * @author Greg Turnquist
  * @author Yuriy Tsarkov
+ * @author Ariel Morelli Andres, Atlassian US, Inc
  */
 public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, QueryComment {
 
@@ -64,7 +68,7 @@ public enum PersistenceProvider implements QueryExtractor, ProxyIdAccessor, Quer
 	 * @see <a href="https://github.com/spring-projects/spring-data-jpa/issues/846">DATAJPA-444</a>
 	 */
 	HIBERNATE(//
-			Collections.singletonList(HIBERNATE_ENTITY_MANAGER_INTERFACE), //
+			Collections.singletonList(HIBERNATE_ENTITY_MANAGER_FACTORY_INTERFACE), //
 			Collections.singletonList(HIBERNATE_JPA_METAMODEL_TYPE)) {
 
 		@Override
@@ -114,7 +118,7 @@ public String getCommentHintKey() {
 	/**
 	 * EclipseLink persistence provider.
 	 */
-	ECLIPSELINK(Collections.singleton(ECLIPSELINK_ENTITY_MANAGER_INTERFACE),
+	ECLIPSELINK(Collections.singleton(ECLIPSELINK_ENTITY_MANAGER_FACTORY_INTERFACE),
 			Collections.singleton(ECLIPSELINK_JPA_METAMODEL_TYPE)) {
 
 		@Override
@@ -152,7 +156,7 @@ public String getCommentHintValue(String comment) {
 	/**
 	 * Unknown special provider. Use standard JPA.
 	 */
-	GENERIC_JPA(Collections.singleton(GENERIC_JPA_ENTITY_MANAGER_INTERFACE), Collections.emptySet()) {
+	GENERIC_JPA(Collections.singleton(GENERIC_JPA_ENTITY_MANAGER_FACTORY_INTERFACE), Collections.emptySet()) {
 
 		@Nullable
 		@Override
@@ -199,25 +203,25 @@ public String getCommentHintKey() {
 	private static final Collection<PersistenceProvider> ALL = List.of(HIBERNATE, ECLIPSELINK, GENERIC_JPA);
 
 	private static final ConcurrentReferenceHashMap<Class<?>, PersistenceProvider> CACHE = new ConcurrentReferenceHashMap<>();
-	private final Iterable<String> entityManagerClassNames;
+	private final Iterable<String> entityManagerFactoryClassNames;
 	private final Iterable<String> metamodelClassNames;
 
 	private final boolean present;
 
 	/**
 	 * Creates a new {@link PersistenceProvider}.
 	 *
-	 * @param entityManagerClassNames the names of the provider specific {@link EntityManager} implementations. Must not
-	 *          be {@literal null} or empty.
+	 * @param entityManagerFactoryClassNames the names of the provider specific {@link EntityManagerFactory}
+	 *          implementations. Must not be {@literal null} or empty.
 	 * @param metamodelClassNames must not be {@literal null}.
 	 */
-	PersistenceProvider(Iterable<String> entityManagerClassNames, Iterable<String> metamodelClassNames) {
+	PersistenceProvider(Iterable<String> entityManagerFactoryClassNames, Iterable<String> metamodelClassNames) {
 
-		this.entityManagerClassNames = entityManagerClassNames;
+		this.entityManagerFactoryClassNames = entityManagerFactoryClassNames;
 		this.metamodelClassNames = metamodelClassNames;
 
 		boolean present = false;
-		for (String entityManagerClassName : entityManagerClassNames) {
+		for (String entityManagerClassName : this.entityManagerFactoryClassNames) {
 
 			if (ClassUtils.isPresent(entityManagerClassName, PersistenceProvider.class.getClassLoader())) {
 				present = true;
@@ -250,23 +254,26 @@ private static PersistenceProvider cacheAndReturn(Class<?> type, PersistenceProv
 	public static PersistenceProvider fromEntityManager(EntityManager em) {
 
 		Assert.notNull(em, "EntityManager must not be null");
-
-		Class<?> entityManagerType = em.getDelegate().getClass();
-		PersistenceProvider cachedProvider = CACHE.get(entityManagerType);
+		EntityManagerFactory entityManagerFactory = em.getEntityManagerFactory();
+		EntityManagerFactory targetEntityManagerFactory = Proxy.isProxyClass(entityManagerFactory.getClass())
+				? Objects.requireNonNullElse(entityManagerFactory.unwrap(null), entityManagerFactory)
+				: entityManagerFactory;
+		Class<? extends EntityManagerFactory> entityManagerFactoryType = targetEntityManagerFactory.getClass();
+		PersistenceProvider cachedProvider = CACHE.get(entityManagerFactoryType);
 
 		if (cachedProvider != null) {
 			return cachedProvider;
 		}
 
 		for (PersistenceProvider provider : ALL) {
-			for (String entityManagerClassName : provider.entityManagerClassNames) {
-				if (isEntityManagerOfType(em, entityManagerClassName)) {
-					return cacheAndReturn(entityManagerType, provider);
+			for (String entityManagerClassName : provider.entityManagerFactoryClassNames) {
+				if (isEntityManagerFactoryOfType(targetEntityManagerFactory, entityManagerClassName)) {
+					return cacheAndReturn(entityManagerFactoryType, provider);
 				}
 			}
 		}
 
-		return cacheAndReturn(entityManagerType, GENERIC_JPA);
+		return cacheAndReturn(entityManagerFactoryType, GENERIC_JPA);
 	}
 
 	/**
@@ -354,10 +361,10 @@ public boolean isPresent() {
 	 */
 	interface Constants {
 
-		String GENERIC_JPA_ENTITY_MANAGER_INTERFACE = "jakarta.persistence.EntityManager";
-		String ECLIPSELINK_ENTITY_MANAGER_INTERFACE = "org.eclipse.persistence.jpa.JpaEntityManager";
+		String GENERIC_JPA_ENTITY_MANAGER_FACTORY_INTERFACE = "jakarta.persistence.EntityManagerFactory";
+		String ECLIPSELINK_ENTITY_MANAGER_FACTORY_INTERFACE = "org.eclipse.persistence.jpa.JpaEntityManagerFactory";
 		// needed as Spring only exposes that interface via the EM proxy
-		String HIBERNATE_ENTITY_MANAGER_INTERFACE = "org.hibernate.engine.spi.SessionImplementor";
+		String HIBERNATE_ENTITY_MANAGER_FACTORY_INTERFACE = "org.hibernate.SessionFactory";
 
 		String HIBERNATE_JPA_METAMODEL_TYPE = "org.hibernate.metamodel.model.domain.JpaMetamodel";
 		String ECLIPSELINK_JPA_METAMODEL_TYPE = "org.eclipse.persistence.internal.jpa.metamodel.MetamodelImpl";

spring-data-jpa/src/test/java/org/springframework/data/jpa/provider/PersistenceProviderUnitTests.java

@@ -21,9 +21,11 @@
 
 import jakarta.persistence.EntityManager;
 
+import java.lang.reflect.Proxy;
 import java.util.Arrays;
 import java.util.Map;
 
+import jakarta.persistence.EntityManagerFactory;
 import org.assertj.core.api.Assumptions;
 import org.hibernate.Version;
 import org.junit.jupiter.api.BeforeEach;
@@ -42,6 +44,7 @@
  * @author Thomas Darimont
  * @author Oliver Gierke
  * @author Jens Schauder
+ * @author Ariel Morelli Andres, Atlassian US, Inc
  */
 class PersistenceProviderUnitTests {
 
@@ -61,15 +64,15 @@ void detectsEclipseLinkPersistenceProvider() throws Exception {
 
 		shadowingClassLoader.excludePackage("org.eclipse.persistence.jpa");
 
-		EntityManager em = mockProviderSpecificEntityManagerInterface(ECLIPSELINK_ENTITY_MANAGER_INTERFACE);
+		EntityManager em = mockEntityManagerWithProviderFactoryInterface(ECLIPSELINK_ENTITY_MANAGER_FACTORY_INTERFACE);
 
 		assertThat(fromEntityManager(em)).isEqualTo(ECLIPSELINK);
 	}
 
 	@Test
 	void fallbackToGenericJpaForUnknownPersistenceProvider() throws Exception {
 
-		EntityManager em = mockProviderSpecificEntityManagerInterface("foo.bar.unknown.jpa.JpaEntityManager");
+		EntityManager em = mockEntityManagerWithProviderFactoryInterface("foo.bar.unknown.jpa.JpaEntityManager");
 
 		assertThat(fromEntityManager(em)).isEqualTo(GENERIC_JPA);
 	}
@@ -81,33 +84,37 @@ void detectsHibernatePersistenceProviderForHibernateVersion52() throws Exception
 
 		shadowingClassLoader.excludePackage("org.hibernate");
 
-		EntityManager em = mockProviderSpecificEntityManagerInterface(HIBERNATE_ENTITY_MANAGER_INTERFACE);
+		EntityManager em = mockEntityManagerWithProviderFactoryInterface(HIBERNATE_ENTITY_MANAGER_FACTORY_INTERFACE);
 
 		assertThat(fromEntityManager(em)).isEqualTo(HIBERNATE);
 	}
 
-	@Test // DATAJPA-1379
-	void detectsProviderFromProxiedEntityManager() throws Exception {
+	@Test
+	void detectsProviderFromProxiedEntityManagerFactory() throws Exception {
 
 		shadowingClassLoader.excludePackage("org.eclipse.persistence.jpa");
 
-		EntityManager em = mockProviderSpecificEntityManagerInterface(ECLIPSELINK_ENTITY_MANAGER_INTERFACE);
-
+		EntityManager em = mockEntityManagerWithProviderFactoryInterface(ECLIPSELINK_ENTITY_MANAGER_FACTORY_INTERFACE);
+		EntityManagerFactory proxiedFactory = (EntityManagerFactory) Proxy.newProxyInstance(getClass().getClassLoader(),
+				em.getEntityManagerFactory().getClass().getInterfaces(), (proxy, method, args) -> switch (method.getName()) {
+					case "unwrap" -> args[0] == null ? em.getEntityManagerFactory() : this;
+					default -> method.invoke(em.getEntityManagerFactory(), args);
+				});
 		EntityManager emProxy = Mockito.mock(EntityManager.class);
-		Mockito.when(emProxy.getDelegate()).thenReturn(em);
+		Mockito.when(emProxy.getEntityManagerFactory()).thenReturn(proxiedFactory);
 
 		assertThat(fromEntityManager(emProxy)).isEqualTo(ECLIPSELINK);
 	}
 
-	private EntityManager mockProviderSpecificEntityManagerInterface(String interfaceName) throws ClassNotFoundException {
+	private EntityManager mockEntityManagerWithProviderFactoryInterface(String factoryInterfaceName)
+			throws ClassNotFoundException {
 
-		Class<?> providerSpecificEntityManagerInterface = InterfaceGenerator.generate(interfaceName, shadowingClassLoader,
-				EntityManager.class);
+		Class<?> providerSpecificEntityManagerFactoryInterface = InterfaceGenerator.generate(factoryInterfaceName,
+				shadowingClassLoader, EntityManagerFactory.class);
 
-		EntityManager em = (EntityManager) Mockito.mock(providerSpecificEntityManagerInterface);
-		Mockito.when(em.getDelegate()).thenReturn(em); // delegate is used to determine the classloader of the provider
-																										// specific interface, therefore we return the proxied
-																										// EntityManager.
+		EntityManagerFactory factory = (EntityManagerFactory) Mockito.mock(providerSpecificEntityManagerFactoryInterface);
+		EntityManager em = Mockito.mock(EntityManager.class);
+		Mockito.when(em.getEntityManagerFactory()).thenReturn(factory);
 
 		return em;
 	}

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/query/AbstractStringBasedJpaQueryUnitTests.java

@@ -18,6 +18,7 @@
 import static org.mockito.Mockito.*;
 
 import jakarta.persistence.EntityManager;
+import jakarta.persistence.EntityManagerFactory;
 import jakarta.persistence.metamodel.Metamodel;
 
 import java.lang.reflect.Method;
@@ -53,6 +54,7 @@
  *
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Ariel Morelli Andres, Atlassian US, Inc
  */
 class AbstractStringBasedJpaQueryUnitTests {
 
@@ -135,15 +137,16 @@ static class InvocationCapturingStringQueryStub extends AbstractStringBasedJpaQu
 				public EntityManager get() {
 
 					EntityManager em = Mockito.mock(EntityManager.class);
+					EntityManagerFactory factory = Mockito.mock(EntityManagerFactory.class);
 
 					Metamodel meta = mock(Metamodel.class);
 					when(em.getMetamodel()).thenReturn(meta);
 					when(em.getDelegate()).thenReturn(new Object()); // some generic jpa
+					when(em.getEntityManagerFactory()).thenReturn(factory); // some generic jpa
 
 					return em;
 				}
-			}.get(), queryString, countQueryString, Mockito.mock(QueryRewriter.class),
-					ValueExpressionDelegate.create());
+			}.get(), queryString, countQueryString, Mockito.mock(QueryRewriter.class), ValueExpressionDelegate.create());
 
 			this.targetMethod = targetMethod;
 		}

spring-data-jpa/src/test/java/org/springframework/data/jpa/repository/support/SimpleJpaRepositoryUnitTests.java

@@ -60,6 +60,7 @@
  * @author Jens Schauder
  * @author Greg Turnquist
  * @author Yanming Zhou
+ * @author Ariel Morelli Andres, Atlassian US, Inc
  */
 @ExtendWith(MockitoExtension.class)
 @MockitoSettings(strictness = Strictness.LENIENT)
@@ -84,6 +85,7 @@ class SimpleJpaRepositoryUnitTests {
 	void setUp() {
 
 		when(em.getDelegate()).thenReturn(em);
+		when(em.getEntityManagerFactory()).thenReturn(entityManagerFactory);
 
 		when(information.getJavaType()).thenReturn(User.class);
 		when(em.getCriteriaBuilder()).thenReturn(builder);