Multi-tenancy using Spring Boot and Hibernate

Implement Multi-tenancy using Spring boot

Chirag S P

2021-11-13

In modern Sass applications, a single application can be used by multiple companies(tenant). There are three ways with which one can divide an application by,

Single schema and a single DB instance
Separate Schema for each company but single DB instance
Different DB instances for each Company

In this article, we will be looking into the second method , i.e Separate Schema for each tenant or company using Spring Boot.

Code Setup

First we need to write the TenantContext class. Here we set the current tenant(schema) in the currently running Thread. The class is as follows,

public class TenantContext {

	private static ThreadLocal<String> currentTenant = new InheritableThreadLocal<>();

	public static void setCurrentTenant(String tenant) {
		currentTenant.set(tenant);
	}

	public static String getCurrentTenant() {
		return currentTenant.get();
	}

	public static void clear() {
		currentTenant.remove();
	}

}

Next we need to implement the ‘CurrentTenantIdentifierResolver’. This is used to set/resolve the current tenant. This class is as follows,

import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
import org.springframework.stereotype.Component;

import com.test.www.test.tenant.TenantContext;
import com.test.www.test.util.OotaConstants;

@Component
public class CurrentTenantIdentifierResolverImpl implements CurrentTenantIdentifierResolver {

	 private String defaultTenant ="DEFAULT_TENANT";

	    @Override
	    public String resolveCurrentTenantIdentifier() {
	        String t =  TenantContext.getCurrentTenant();
	        if(t!=null){
	            return t;
	        } else {
	            return defaultTenant;
	        }
	    }

	    @Override
	    public boolean validateExistingCurrentSessions() {
	        return true;
	    }
	}

The next step is to define ‘MultiTenantConnectionProvider’ class. Here we define how we go about getting the connection to Database. This class is as follows,

import java.sql.Connection;
import java.sql.SQLException;

import javax.sql.DataSource;

import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
public class MultiTenantConnectionProviderImpl implements MultiTenantConnectionProvider {

	/**
	 * 
	 */
	private static final long serialVersionUID = 1L;

	@Autowired
	private DataSource dataSource;
	private static Logger logger = LoggerFactory.getLogger(MultiTenantConnectionProviderImpl.class);

	public MultiTenantConnectionProviderImpl(DataSource dataSource) {
		this.dataSource = dataSource;
	}

	@Override
	public Connection getAnyConnection() throws SQLException {
		return dataSource.getConnection();
	}

	@Override
	public void releaseAnyConnection(Connection connection) throws SQLException {
		connection.close();
	}

	@Override
	public Connection getConnection(String tenantIdentifier) throws SQLException {
		logger.info("Get connection for tenant {}", tenantIdentifier);
		final Connection connection = getAnyConnection();
		connection.setSchema(tenantIdentifier);
		return connection;
	}

	@Override
	public void releaseConnection(String tenantIdentifier, Connection connection) throws SQLException {
		logger.info("Release connection for tenant {}", tenantIdentifier);
		connection.setSchema("DEFAULT_TENANT");
		releaseAnyConnection(connection);

	}

	@Override
	public boolean supportsAggressiveRelease() {
		return false;
	}

	@Override
	@SuppressWarnings("rawtypes")
	public boolean isUnwrappableAs(Class unwrapType) {
		return false;
	}

	@Override
	public <T> T unwrap(Class<T> unwrapType) {
		return null;
	}

}

Add interceptor to get the tenant(schema) details from request as below,

@Component
public class TenantInterceptor extends HandlerInterceptorAdapter {

  private static final String TENANT_HEADER = "TenantID";

  @Override
  public boolean preHandle(HttpServletRequest req, HttpServletResponse res, Object handler)
      throws Exception {

    String tenant = req.getHeader(TENANT_HEADER);
    boolean tenantSet = false;

    if(StringUtils.isEmpty(tenant)) {
      res.setStatus(HttpServletResponse.SC_BAD_REQUEST);
      res.setContentType(MediaType.APPLICATION_JSON_VALUE);
      res.getWriter().write("{\"error\": \"No tenant supplied\"}");
      res.getWriter().flush();
    } else {
      TenantContext.setCurrentTenant(tenant);
      tenantSet = true;
    }

    return tenantSet;
  }

  @Override
  public void postHandle(
      HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView)
          throws Exception {
    TenantContext.clear();
  }

Here TenantID is the header parameter. The tenant(schema) which we desire to connect to should be supplied here.

Finally we need to set the hibernate configuration. We need to inform hibernate that we are using multi-tenancy and which kind of multi-tenancy at that. Code for the same is as below,

import java.util.HashMap;
import java.util.Map;

import javax.sql.DataSource;

import org.hibernate.MultiTenancyStrategy;
import org.hibernate.cfg.Environment;
import org.hibernate.context.spi.CurrentTenantIdentifierResolver;
import org.hibernate.engine.jdbc.connections.spi.MultiTenantConnectionProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.orm.jpa.JpaProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.orm.jpa.JpaVendorAdapter;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.orm.jpa.vendor.HibernateJpaVendorAdapter;

@Configuration
public class HibernateConfig {

  @Autowired
  private JpaProperties jpaProperties;

  @Bean
  public JpaVendorAdapter jpaVendorAdapter() {
    return new HibernateJpaVendorAdapter();
  }

  @Bean
  public LocalContainerEntityManagerFactoryBean entityManagerFactory(DataSource dataSource,
      MultiTenantConnectionProvider multiTenantConnectionProviderImpl,
      CurrentTenantIdentifierResolver currentTenantIdentifierResolverImpl) {
    Map<String, Object> properties = new HashMap<>();
    properties.putAll(jpaProperties.getProperties());
    properties.put(Environment.MULTI_TENANT, MultiTenancyStrategy.SCHEMA);
    properties.put(Environment.MULTI_TENANT_CONNECTION_PROVIDER, multiTenantConnectionProviderImpl);
    properties.put(Environment.MULTI_TENANT_IDENTIFIER_RESOLVER, currentTenantIdentifierResolverImpl);

    LocalContainerEntityManagerFactoryBean em = new LocalContainerEntityManagerFactoryBean();
    em.setDataSource(dataSource);
    em.setPackagesToScan("com.test.www.test");
    em.setJpaVendorAdapter(jpaVendorAdapter());
    em.setJpaPropertyMap(properties);
    return em;
  }
}

In the code shown above, there are some configurations which need to be set for multi-tenancy to work,

properties.put(Environment.MULTI_TENANT, MultiTenancyStrategy.SCHEMA);
properties.put(Environment.MULTI_TENANT_CONNECTION_PROVIDER, multiTenantConnectionProviderImpl);
properties.put(Environment.MULTI_TENANT_IDENTIFIER_RESOLVER, currentTenantIdentifierResolverImpl);

Testing

Time to test multi-tenancy in a simple REST application

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RestController;

import com.test.www.test.pojos.User;
import com.test.www.test.services.UserService;

@RestController
public class UserController {
	
	@Autowired
	private UserService userRegistrationService;
	
	@GetMapping("/user/{userName}")
	public User getUserDetailsByUserName(@PathVariable(name = "userName",required=true) String userName) {
		return userRegistrationService.getUserByUserName(userName);
	}
}

package com.test.www.test.services.impl;

import java.util.Date;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import com.test.www.test.entities.UserDetails;
import com.test.www.test.repository.UserDetailsRepository;
import com.test.www.test.services.UserService;

@Service
public class UserServiceImpl implements UserService{
	@Autowired
	private UserDetailsRepository userDetailsRepository;
	
	@Override
	public UserDetails getUserByUserName(String username) {
		UserDetails userDetails = userDetailsRepository.getUserByUserName(username);
		return userDetails;
	}
}

The repository class ‘UserDetailsRepository’ is as shown below,

package com.test.www.test.repository;

import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param;
import org.springframework.stereotype.Repository;

import com.test.www.test.entities.UserDetails;

@Repository
public interface UserDetailsRepository extends JpaRepository<UserDetails,Integer>{
	@Query(value="SELECT u FROM UserDetails u where userName =:user_name")
	UserDetails getUserByUserName(@Param("user_name")String userName);
}

So that is how we go about implementing multi-tenancy in Java using Spring Boot and Hibernate.**