HHH-7312 - Fixed missing usage of CurrentTenantIdentifierResolver in EntityManagerImpl

hibernate-entitymanager/src/main/java/org/hibernate/ejb/EntityManagerImpl.java

@@ -37,6 +37,7 @@
 import org.hibernate.annotations.common.util.ReflectHelper;
 import org.hibernate.ejb.internal.EntityManagerMessageLogger;
 import org.hibernate.engine.spi.SessionBuilderImplementor;
+import org.hibernate.engine.spi.SessionFactoryImplementor;
 import org.hibernate.engine.spi.SessionImplementor;
 
 /**
@@ -100,8 +101,14 @@ public Session getSession() {
 	@Override
     protected Session getRawSession() {
 		if ( session == null ) {
-			SessionBuilderImplementor sessionBuilder = getEntityManagerFactory().getSessionFactory().withOptions();
+			SessionFactoryImplementor sessionFactory = (SessionFactoryImplementor) getEntityManagerFactory().getSessionFactory();
+			SessionBuilderImplementor sessionBuilder = sessionFactory.withOptions();
 			sessionBuilder.owner( this );
+
+			if (sessionFactory.getCurrentTenantIdentifierResolver()!=null){
+				sessionBuilder.tenantIdentifier(sessionFactory.getCurrentTenantIdentifierResolver().resolveCurrentTenantIdentifier());
+			}
+
 			if (sessionInterceptorClass != null) {
 				try {
 					Interceptor interceptor = (Interceptor) sessionInterceptorClass.newInstance();

hibernate-entitymanager/src/test/java/org/hibernate/ejb/test/multitenancy/CurrentTenantResolverTest.java

@@ -0,0 +1,70 @@
+package org.hibernate.ejb.test.multitenancy;
+
+import static org.junit.Assert.*;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.util.Map;
+import java.util.Properties;
+
+import javax.persistence.EntityManager;
+
+import junit.framework.Assert;
+
+import org.hibernate.Session;
+import org.hibernate.cfg.AvailableSettings;
+import org.hibernate.cfg.Configuration;
+import org.hibernate.dialect.Dialect;
+import org.hibernate.dialect.H2Dialect;
+import org.hibernate.ejb.Ejb3Configuration;
+import org.hibernate.ejb.EntityManagerImpl;
+import org.hibernate.ejb.test.BaseEntityManagerFunctionalTestCase;
+import org.hibernate.testing.TestForIssue;
+import org.hibernate.testing.env.ConnectionProviderBuilder;
+import org.junit.Test;
+
+@TestForIssue(jiraKey="HHH-7312")
+public class CurrentTenantResolverTest extends
+		BaseEntityManagerFunctionalTestCase {
+
+	private Properties properties;
+	private TestCurrentTenantIdentifierResolver tenantResolver;
+
+	@Override
+	protected Dialect getDialect() {
+		return H2Dialect.getDialect();
+	}
+
+	@Override
+	protected Ejb3Configuration constructConfiguration() {		
+		Ejb3Configuration ejb3cfg=super.constructConfiguration();
+		Configuration cfg=ejb3cfg.getHibernateConfiguration();
+
+		properties = new Properties();
+		properties.put(org.hibernate.ejb.AvailableSettings.JDBC_DRIVER,
+				"org.h2.Driver");
+		properties.put(org.hibernate.ejb.AvailableSettings.JDBC_URL,
+				"jdbc:h2:mem:");
+		properties.put(AvailableSettings.MULTI_TENANT, "SCHEMA");
+		tenantResolver = new TestCurrentTenantIdentifierResolver();
+		properties.put(AvailableSettings.MULTI_TENANT_IDENTIFIER_RESOLVER,tenantResolver);
+
+		TestMultiTenantConnectionProvider provider=new TestMultiTenantConnectionProvider();
+		provider.addProvider("acme",ConnectionProviderBuilder.buildConnectionProvider("acme"));
+		provider.addProvider("jboss",ConnectionProviderBuilder.buildConnectionProvider("jboss"));
+
+		properties.put(AvailableSettings.MULTI_TENANT_CONNECTION_PROVIDER,provider);
+
+		cfg.setProperties(properties);
+		return ejb3cfg;
+	}
+
+	@Test
+	public void testCurrentTenantIdentifierResolverIsUsed() {
+		tenantResolver.setCurrentTenantIdentifier("acme");
+		EntityManager em=createEntityManager(properties);
+		String usedTenant=((EntityManagerImpl)em).getSession().getTenantIdentifier();
+		Assert.assertEquals("acme", usedTenant);
+		em.close();
+	}
+}

hibernate-entitymanager/src/test/java/org/hibernate/ejb/test/multitenancy/TestCurrentTenantIdentifierResolver.java

@@ -0,0 +1,22 @@
+package org.hibernate.ejb.test.multitenancy;
+
+import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
+
+public class TestCurrentTenantIdentifierResolver implements
+		CurrentTenantIdentifierResolver {
+	private String currentTenantIdentifier;
+
+	@Override
+	public boolean validateExistingCurrentSessions() {
+		return false;
+	}
+
+	@Override
+	public String resolveCurrentTenantIdentifier() {
+		return currentTenantIdentifier;
+	}
+
+	public void setCurrentTenantIdentifier(String tenantId) {
+		this.currentTenantIdentifier = tenantId;
+	}
+}

hibernate-entitymanager/src/test/java/org/hibernate/ejb/test/multitenancy/TestMultiTenantConnectionProvider.java

@@ -0,0 +1,32 @@
+package org.hibernate.ejb.test.multitenancy;
+
+import java.util.HashMap;
+
+import org.hibernate.HibernateException;
+import org.hibernate.service.jdbc.connections.internal.DriverManagerConnectionProviderImpl;
+import org.hibernate.service.jdbc.connections.spi.AbstractMultiTenantConnectionProvider;
+import org.hibernate.service.jdbc.connections.spi.ConnectionProvider;
+
+public class TestMultiTenantConnectionProvider extends AbstractMultiTenantConnectionProvider{
+
+	private HashMap<String, ConnectionProvider> providerMap=new HashMap<String, ConnectionProvider>();
+
+	@Override
+	protected ConnectionProvider getAnyConnectionProvider() {
+		return providerMap.entrySet().iterator().next().getValue();
+	}
+
+	@Override
+	protected ConnectionProvider selectConnectionProvider(String tenantIdentifier) {
+		ConnectionProvider provider=providerMap.get(tenantIdentifier);
+		if (provider==null)
+			throw new HibernateException( "Unknown tenant identifier" );
+		return provider;
+	}
+
+	public void addProvider(String tenantIdentifier,
+			DriverManagerConnectionProviderImpl connectionProvider) {
+		providerMap.put(tenantIdentifier, connectionProvider);
+
+	}
+}