HHH-15422 Pick up CurrentTenantIdentifierResolver and MultiTenantConnectionProvider from BeanContainer if not explicit set

hibernate-core/src/main/java/org/hibernate/boot/internal/SessionFactoryOptionsBuilder.java

@@ -63,6 +63,10 @@ import org.hibernate.query.sqm.function.SqmFunctionRegistry;
 import org.hibernate.query.sqm.mutation.spi.SqmMultiTableInsertStrategy;
 import org.hibernate.query.sqm.mutation.spi.SqmMultiTableMutationStrategy;
 import org.hibernate.query.sqm.sql.SqmTranslatorFactory;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.internal.Helper;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
 import org.hibernate.resource.jdbc.spi.PhysicalConnectionHandlingMode;
 import org.hibernate.resource.jdbc.spi.StatementInspector;
 import org.hibernate.resource.transaction.spi.TransactionCoordinatorBuilder;
@@ -283,6 +287,7 @@ public class SessionFactoryOptionsBuilder implements SessionFactoryOptions {
 	private final int queryStatisticsMaxSize;
 
 
+	@SuppressWarnings( "unchecked" )
 	public SessionFactoryOptionsBuilder(StandardServiceRegistry serviceRegistry, BootstrapContext context) {
 		this.serviceRegistry = serviceRegistry;
 		this.jpaBootstrap = context.isJpaBootstrap();
@@ -372,6 +377,38 @@ public class SessionFactoryOptionsBuilder implements SessionFactoryOptions {
 				CurrentTenantIdentifierResolver.class,
 				configurationSettings.get( MULTI_TENANT_IDENTIFIER_RESOLVER )
 		);
+		if ( this.currentTenantIdentifierResolver == null ) {
+			final BeanContainer beanContainer = Helper.allowExtensionsInCdi( serviceRegistry ) ? serviceRegistry.requireService( ManagedBeanRegistry.class ).getBeanContainer() : null;
+			if (beanContainer != null) {
+				this.currentTenantIdentifierResolver = beanContainer.getBean(
+						CurrentTenantIdentifierResolver.class,
+						new BeanContainer.LifecycleOptions() {
+							@Override
+							public boolean canUseCachedReferences() {
+								return true;
+							}
+
+							@Override
+							public boolean useJpaCompliantCreation() {
+								return false;
+							}
+						},
+						new BeanInstanceProducer() {
+
+							@Override
+							public <B> B produceBeanInstance(Class<B> beanType) {
+								return null;
+							}
+
+							@Override
+							public <B> B produceBeanInstance(String name, Class<B> beanType) {
+								return null;
+							}
+
+						}
+				).getBeanInstance();
+			}
+		}
 
 		this.delayBatchFetchLoaderCreations = configurationService.getSetting( DELAY_ENTITY_LOADER_CREATIONS, BOOLEAN, true );
 		this.defaultBatchFetchSize = getInt( DEFAULT_BATCH_FETCH_SIZE, configurationSettings, -1 );

hibernate-core/src/main/java/org/hibernate/engine/jdbc/connections/internal/MultiTenantConnectionProviderInitiator.java

@@ -12,6 +12,10 @@ import org.hibernate.boot.registry.classloading.spi.ClassLoadingException;
 import org.hibernate.cfg.AvailableSettings;
 import org.hibernate.engine.jdbc.connections.spi.DataSourceBasedMultiTenantConnectionProviderImpl;
 import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.internal.Helper;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+import org.hibernate.resource.beans.spi.ManagedBeanRegistry;
 import org.hibernate.service.spi.ServiceException;
 import org.hibernate.service.spi.ServiceRegistryImplementor;
 
@@ -39,7 +43,36 @@ public class MultiTenantConnectionProviderInitiator implements StandardServiceIn
 	@Override
 	public MultiTenantConnectionProvider<?> initiateService(Map<String, Object> configurationValues, ServiceRegistryImplementor registry) {
 		if ( !configurationValues.containsKey( AvailableSettings.MULTI_TENANT_CONNECTION_PROVIDER ) ) {
-			// nothing to do, but given the separate hierarchies have to handle this here.
+			final BeanContainer beanContainer = Helper.allowExtensionsInCdi( registry ) ? registry.requireService( ManagedBeanRegistry.class ).getBeanContainer() : null;
+			if (beanContainer != null) {
+				return beanContainer.getBean(
+						MultiTenantConnectionProvider.class,
+						new BeanContainer.LifecycleOptions() {
+							@Override
+							public boolean canUseCachedReferences() {
+								return true;
+							}
+
+							@Override
+							public boolean useJpaCompliantCreation() {
+								return true;
+							}
+						},
+						new BeanInstanceProducer() {
+
+							@Override
+							public <B> B produceBeanInstance(Class<B> beanType) {
+								return null;
+							}
+
+							@Override
+							public <B> B produceBeanInstance(String name, Class<B> beanType) {
+								return null;
+							}
+
+						}
+				).getBeanInstance();
+			}
 			return null;
 		}
 

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/SimpleBeanContainer.java

@@ -11,7 +11,7 @@ import org.hibernate.resource.beans.container.spi.ContainedBean;
 import org.hibernate.resource.beans.spi.BeanInstanceProducer;
 
 /**
- *
+ * @author Yanming Zhou
  */
 @SuppressWarnings("unchecked")
 public class SimpleBeanContainer implements BeanContainer {
@@ -23,10 +23,8 @@ public class SimpleBeanContainer implements BeanContainer {
 			Class<B> beanType,
 			LifecycleOptions lifecycleOptions,
 			BeanInstanceProducer fallbackProducer) {
-		if ( beanType == SimpleGenerator.class ) {
+		return () -> (B) ( beanType == SimpleGenerator.class ?
-			return () -> (B) new SimpleGenerator( new AtomicLong( INITIAL_VALUE ) );
+				new SimpleGenerator( new AtomicLong( INITIAL_VALUE ) ) : fallbackProducer.produceBeanInstance( beanType ) );
-		}
-		return null;
 	}
 
 	@Override

hibernate-core/src/test/java/org/hibernate/orm/test/idgen/userdefined/UserDefinedGeneratorsTests.java

@@ -24,6 +24,7 @@ import org.hibernate.resource.beans.container.spi.BeanContainer;
 import org.hibernate.resource.beans.container.spi.BeanContainer.LifecycleOptions;
 import org.hibernate.resource.beans.container.spi.ContainedBean;
 import org.hibernate.resource.beans.internal.FallbackBeanInstanceProducer;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
 
 import org.hibernate.testing.orm.junit.JiraKey;
 import org.hibernate.testing.orm.junit.BaseUnitTest;
@@ -39,7 +40,7 @@ import org.mockito.Mockito;
 
 import static org.hamcrest.CoreMatchers.instanceOf;
 import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertThat;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.same;
 import static org.mockito.BDDMockito.given;
@@ -59,10 +60,16 @@ public class UserDefinedGeneratorsTests {
 
 		final BeanContainer beanContainer = Mockito.mock( BeanContainer.class );
 		given(beanContainer.getBean( any(), any(), any() ) ).willAnswer( invocation -> {
+			Class<?> beanType = (Class<?>) invocation.getArguments()[0];
 			LifecycleOptions options = (LifecycleOptions) invocation.getArguments()[1];
-			assertThat( options.canUseCachedReferences(), is( false ) );
+			if (beanType == TestIdentifierGenerator.class) {
-			assertThat( options.useJpaCompliantCreation(), is( true ) );
+				assertThat( options.canUseCachedReferences(), is( false ) );
-			return (ContainedBean<?>) TestIdentifierGenerator::new;
+				assertThat( options.useJpaCompliantCreation(), is( true ) );
+				return (ContainedBean<?>) TestIdentifierGenerator::new;
+			}
+			else {
+				return (ContainedBean<?>) () -> ( ( BeanInstanceProducer ) invocation.getArguments()[2] ).produceBeanInstance( beanType );
+			}
 		} );
 
 		final StandardServiceRegistryBuilder ssrb = ServiceRegistryUtil.serviceRegistryBuilder();

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/AbstractMultiTenancyTest.java

@@ -55,9 +55,9 @@ public abstract class AbstractMultiTenancyTest extends BaseUnitTestCase {
 	protected static final String FRONT_END_TENANT = "front_end";
 	protected static final String BACK_END_TENANT = "back_end";
 
-	private Map<String, ConnectionProvider> connectionProviderMap = new HashMap<>();
+	protected Map<String, ConnectionProvider> connectionProviderMap = new HashMap<>();
 
-	private SessionFactory sessionFactory;
+	protected SessionFactory sessionFactory;
 
 	public AbstractMultiTenancyTest() {
 		init();
@@ -67,13 +67,15 @@ public abstract class AbstractMultiTenancyTest extends BaseUnitTestCase {
 	private void init() {
 		registerConnectionProvider(FRONT_END_TENANT);
 		registerConnectionProvider(BACK_END_TENANT);
+		sessionFactory = sessionFactory(createSettings());
+	}
 
+	protected Map<String, Object> createSettings() {
 		Map<String, Object> settings = new HashMap<>();
 
 		settings.put(AvailableSettings.MULTI_TENANT_CONNECTION_PROVIDER,
-			new ConfigurableMultiTenantConnectionProvider(connectionProviderMap));
+				new ConfigurableMultiTenantConnectionProvider(connectionProviderMap));
-
+		return settings;
-		sessionFactory = sessionFactory(settings);
 	}
 	//end::multitenacy-hibernate-MultiTenantConnectionProvider-example[]
 

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/AbstractTenantResolverBeanContainerTest.java

@@ -0,0 +1,34 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.testing.orm.junit.DomainModel;
+import org.hibernate.testing.orm.junit.JiraKey;
+import org.hibernate.testing.orm.junit.SessionFactory;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author Yanming Zhou
+ */
+@JiraKey("HHH-15422")
+@SessionFactory
+@DomainModel(annotatedClasses = TestEntity.class)
+abstract class AbstractTenantResolverBeanContainerTest {
+
+	@Test
+	void tentantIdShouldBeFilled(SessionFactoryScope scope) {
+		scope.inTransaction( s -> {
+			TestEntity entity = new TestEntity();
+			s.persist( entity );
+			s.flush();
+			assertThat( entity.getTenant(), is( TestCurrentTenantIdentifierResolver.FIXED_TENANT ) );
+		} );
+	}
+
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/MultiTenantConnectionProviderFromBeanContainerTest.java

@@ -0,0 +1,80 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.dialect.H2Dialect;
+import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
+import org.hibernate.internal.SessionFactoryImpl;
+import org.hibernate.orm.test.multitenancy.AbstractMultiTenancyTest;
+import org.hibernate.orm.test.multitenancy.ConfigurableMultiTenantConnectionProvider;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.container.spi.ContainedBean;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+
+import org.hibernate.testing.orm.junit.RequiresDialect;
+import org.junit.Test;
+
+import static org.junit.Assert.assertSame;
+
+/**
+ * @author Yanming Zhou
+ */
+@RequiresDialect(H2Dialect.class)
+public class MultiTenantConnectionProviderFromBeanContainerTest extends AbstractMultiTenancyTest {
+
+	private ConfigurableMultiTenantConnectionProvider providerFromBeanContainer;
+
+	@Override
+	protected Map<String, Object> createSettings() {
+		Map<String, Object> settings = new HashMap<>();
+
+		providerFromBeanContainer = new ConfigurableMultiTenantConnectionProvider( connectionProviderMap);
+		settings.put( AvailableSettings.ALLOW_EXTENSIONS_IN_CDI, "true" );
+		settings.put( AvailableSettings.BEAN_CONTAINER, new BeanContainer() {
+			@Override
+			@SuppressWarnings("unchecked")
+			public <B> ContainedBean<B> getBean(
+					Class<B> beanType,
+					LifecycleOptions lifecycleOptions,
+					BeanInstanceProducer fallbackProducer) {
+				return () -> (B) ( beanType == MultiTenantConnectionProvider.class ? providerFromBeanContainer : fallbackProducer.produceBeanInstance( beanType ) );
+			}
+
+			@Override
+			public <B> ContainedBean<B> getBean(
+					String name,
+					Class<B> beanType,
+					LifecycleOptions lifecycleOptions,
+					BeanInstanceProducer fallbackProducer) {
+				return () -> (B) fallbackProducer.produceBeanInstance( beanType );
+			}
+
+			@Override
+			public void stop() {
+
+			}
+		} );
+		return settings;
+	}
+
+	@Override
+	protected String tenantUrl(String originalUrl, String tenantIdentifier) {
+		return originalUrl.replace("db1", tenantIdentifier);
+	}
+
+	@Test
+	public void testProviderInUse() {
+		MultiTenantConnectionProvider<?> providerInUse = ((SessionFactoryImpl) sessionFactory).getServiceRegistry().getService( MultiTenantConnectionProvider.class );
+		assertSame( providerInUse, expectedProviderInUse());
+	}
+
+	protected MultiTenantConnectionProvider<?> expectedProviderInUse() {
+		return providerFromBeanContainer;
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/MultiTenantConnectionProviderFromSettingsOverBeanContainerTest.java

@@ -0,0 +1,36 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import java.util.Map;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.dialect.H2Dialect;
+import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
+import org.hibernate.orm.test.multitenancy.ConfigurableMultiTenantConnectionProvider;
+
+import org.hibernate.testing.orm.junit.RequiresDialect;
+
+/**
+ * @author Yanming Zhou
+ */
+@RequiresDialect(H2Dialect.class)
+public class MultiTenantConnectionProviderFromSettingsOverBeanContainerTest extends MultiTenantConnectionProviderFromBeanContainerTest {
+
+	private ConfigurableMultiTenantConnectionProvider providerFromSettings;
+
+	@Override
+	protected Map<String, Object> createSettings() {
+		Map<String, Object> settings = super.createSettings();
+		providerFromSettings = new ConfigurableMultiTenantConnectionProvider(connectionProviderMap);
+		settings.put(AvailableSettings.MULTI_TENANT_CONNECTION_PROVIDER, providerFromSettings);
+		return settings;
+	}
+
+	@Override
+	protected MultiTenantConnectionProvider<?> expectedProviderInUse() {
+		return providerFromSettings;
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/TenantResolverFromBeanContainerTest.java

@@ -0,0 +1,36 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author Yanming Zhou
+ */
+@ServiceRegistry(
+		settings = {
+				@Setting(name = AvailableSettings.ALLOW_EXTENSIONS_IN_CDI, value = "true"),
+				@Setting(name = AvailableSettings.BEAN_CONTAINER, value = "org.hibernate.orm.test.multitenancy.beancontainer.TestBeanContainer")
+		}
+)
+public class TenantResolverFromBeanContainerTest extends AbstractTenantResolverBeanContainerTest {
+
+	@Test
+	void tenantResolverFromBeanContainerShouldBeUsed(SessionFactoryScope scope) {
+		CurrentTenantIdentifierResolver<?> tenantResolver = scope.getSessionFactory().getCurrentTenantIdentifierResolver();
+		assertThat(tenantResolver, is(TestCurrentTenantIdentifierResolver.INSTANCE_FOR_BEAN_CONTAINER));
+	}
+
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/TenantResolverFromSettingsOverBeanContainerTest.java

@@ -0,0 +1,39 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+
+import org.hibernate.testing.orm.junit.ServiceRegistry;
+import org.hibernate.testing.orm.junit.SessionFactoryScope;
+import org.hibernate.testing.orm.junit.Setting;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.CoreMatchers.not;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author Yanming Zhou
+ */
+@ServiceRegistry(
+		settings = {
+				@Setting(name = AvailableSettings.ALLOW_EXTENSIONS_IN_CDI, value = "true"),
+				@Setting(name = AvailableSettings.MULTI_TENANT_IDENTIFIER_RESOLVER, value = "org.hibernate.orm.test.multitenancy.beancontainer.TestCurrentTenantIdentifierResolver"),
+				@Setting(name = AvailableSettings.BEAN_CONTAINER, value = "org.hibernate.orm.test.multitenancy.beancontainer.TestBeanContainer")
+		}
+)
+public class TenantResolverFromSettingsOverBeanContainerTest extends AbstractTenantResolverBeanContainerTest {
+
+	@Test
+	void tenantResolverFromSettingsShouldBeUsed(SessionFactoryScope scope) {
+		CurrentTenantIdentifierResolver<?> tenantResolver = scope.getSessionFactory().getCurrentTenantIdentifierResolver();
+		assertThat(tenantResolver, instanceOf(TestCurrentTenantIdentifierResolver.class));
+		assertThat(tenantResolver, is(not(TestCurrentTenantIdentifierResolver.INSTANCE_FOR_BEAN_CONTAINER)));
+	}
+
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/TestBeanContainer.java

@@ -0,0 +1,40 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+import org.hibernate.resource.beans.container.spi.BeanContainer;
+import org.hibernate.resource.beans.container.spi.ContainedBean;
+import org.hibernate.resource.beans.spi.BeanInstanceProducer;
+
+/**
+ * @author Yanming Zhou
+ */
+@SuppressWarnings("unchecked")
+public class TestBeanContainer implements BeanContainer {
+
+	@Override
+	public <B> ContainedBean<B> getBean(
+			Class<B> beanType,
+			LifecycleOptions lifecycleOptions,
+			BeanInstanceProducer fallbackProducer) {
+		return () -> (B) ( beanType == CurrentTenantIdentifierResolver.class ?
+				TestCurrentTenantIdentifierResolver.INSTANCE_FOR_BEAN_CONTAINER : fallbackProducer.produceBeanInstance( beanType ) );
+	}
+
+	@Override
+	public <B> ContainedBean<B> getBean(
+			String name,
+			Class<B> beanType,
+			LifecycleOptions lifecycleOptions,
+			BeanInstanceProducer fallbackProducer) {
+		return null;
+	}
+
+	@Override
+	public void stop() {
+
+	}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/TestCurrentTenantIdentifierResolver.java

@@ -0,0 +1,27 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+
+/**
+ * @author Yanming Zhou
+ */
+public class TestCurrentTenantIdentifierResolver implements CurrentTenantIdentifierResolver<String> {
+
+		public static final String FIXED_TENANT = "FIXED";
+
+		public static final CurrentTenantIdentifierResolver<String> INSTANCE_FOR_BEAN_CONTAINER = new TestCurrentTenantIdentifierResolver();
+
+		@Override
+		public boolean validateExistingCurrentSessions() {
+			return false;
+		}
+
+		@Override
+		public String resolveCurrentTenantIdentifier() {
+			return FIXED_TENANT;
+		}
+}

hibernate-core/src/test/java/org/hibernate/orm/test/multitenancy/beancontainer/TestEntity.java

@@ -0,0 +1,48 @@
+/*
+ * SPDX-License-Identifier: LGPL-2.1-or-later
+ * Copyright Red Hat Inc. and Hibernate Authors
+ */
+package org.hibernate.orm.test.multitenancy.beancontainer;
+
+import org.hibernate.annotations.TenantId;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.GeneratedValue;
+import jakarta.persistence.Id;
+
+/**
+ * @author Yanming Zhou
+ */
+@Entity
+public class TestEntity {
+
+	@Id
+	@GeneratedValue
+	public Long id;
+
+	private String name;
+
+	@TenantId
+	private String tenant;
+
+	public String getName() {
+		return name;
+	}
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	public String getTenant() {
+		return tenant;
+	}
+
+	public void setTenant(String tenant) {
+		this.tenant = tenant;
+	}
+
+	public Long getId() {
+		return id;
+	}
+
+}