Multi-Tenancy Implementation for Spring Boot + Hibernate Projects

In this article, we'll see how to achieve multi-tenancy in a Spring Boot project using a schema-per-tenant approach.

This article will also offer a way to handle logins in a multi-tenant environment. This is tricky because, at that stage, it is not always clear to which tenant the yet-to-be-logged user belongs.

This article assumes that the reader is already familiar with multi-tenancy practices, but just in case, here's a short recap:

Multi-Tenancy Has Three Different Approaches

1. DB per tenant: Each tenant has its own DB for its data. This is the highest level of isolation.

2. Schema per tenant: Each tenant's data is saved on the same DB but on a different schema. This approach can be implemented in two different ways:

   • Connection pool per schema

   • Single connection pool for all schemas — for each request, a connection is retrieved from the pool and set schema is called with the relevant tenant before assigning it to the context.

3. Discriminator field: All tenants' data is saved on the same tables, providing that a discriminator field is available on these tables to distinguish each tenant.

I'm not going to dive into the pros and cons of each approach, but if you want to learn more about it, you can read this article, and this MSDN page.

Implementation

In this article, I chose to implement multitenancy using schema-per-tenant paradigm with one connection pool for all tenants.

To handle login, we'll use a general schema (tenant) that has only a single table that maps each user in the system to its relevant tenant. The purpose of this table is to get the user's tenant identifier upon login where the tenant is still unknown. Then, the tenant identifier will be saved in a JWT, but it can be saved in different places, such as an HTTP header.

Multi-Tenancy Setup

First of all, we'll need a shared context of the current tenant. The tenant will be set before each request is handled and be released after. Also, notice that the context is ThreadLocal rather than static because the server can handle more than one tenant at a time.

public class TenantContext {

    private static Logger logger = LoggerFactory.getLogger(TenantContext.class.getName());

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

    public static void setCurrentTenant(String tenant) {
        logger.debug("Setting tenant to " + tenant);
        currentTenant.set(tenant);
    }

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

    public static void clear() {
        currentTenant.set(null);
    }
}

The next is TenantInterceptor, an interceptor that reads the tenant identifier from the JWT (or request header in different implementations) and sets the tenant context:

@Component
public class TenantInterceptor extends HandlerInterceptorAdapter {

    @Autowired
    private JwtTokenUtil jwtTokenUtil;

    @Value("${jwt.header}")
    private String tokenHeader;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler)
            throws Exception {

        String authToken = request.getHeader(this.tokenHeader);
        String tenantId = jwtTokenUtil.getTenantIdFromToken(authToken);
        TenantContext.setCurrentTenant(tenantId);

        return true;
    }

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

Create a  CurrentTenantIdentifierResolver — this is a module that Hibernate needs to resolve the current tenant:

@Component
public class TenantIdentifierResolver implements CurrentTenantIdentifierResolver {

    @Override
    public String resolveCurrentTenantIdentifier() {
        String tenantId = TenantContext.getCurrentTenant();
        if (tenantId != null) {
            return tenantId;
        }
        return DEFAULT_TENANT_ID;
    }

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

The next is MultiTenantConnectionProvider — also required by Hibernate to provide the connection to the context. In our case, we're asking a connection from the data source and setting its schema to the relevant tenant:

@Component
public class MultiTenantConnectionProviderImpl implements MultiTenantConnectionProvider {

    @Autowired
    private 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 tenantIdentifie) throws SQLException {
        String tenantIdentifier = TenantContext.getCurrentTenant();
        final Connection connection = getAnyConnection();
        try {
            if (tenantIdentifier != null) {
                connection.createStatement().execute("USE " + tenantIdentifier);
            } else {
                connection.createStatement().execute("USE " + DEFAULT_TENANT_ID);
            }
        }
        catch ( SQLException e ) {
            throw new HibernateException(
                    "Problem setting schema to " + tenantIdentifier,
                    e
            );
        }
        return connection;
    }

    @Override
    public void releaseConnection(String tenantIdentifier, Connection connection) throws SQLException {
        try {
            connection.createStatement().execute( "USE " + DEFAULT_TENANT_ID );
        }
        catch ( SQLException e ) {
            throw new HibernateException(
                    "Problem setting schema to " + tenantIdentifier,
                    e
            );
        }
        connection.close();
    }

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

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

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

Now to wire it up:

@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.getHibernateProperties(dataSource));
        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.autorni");
        em.setJpaVendorAdapter(jpaVendorAdapter());
        em.setJpaPropertyMap(properties);
        return em;
    }
}

Login Handling

Upon login, we'll need to query the general schema and to retrieve the user's tenantId. Only then do we proceed with the login on the relevant tenant schema.

Notice that once the connection is established in the context (typically when the first query is executed), it is cached per thread and it cannot be changed. Therefore, the tenant cannot be altered in the middle of the controller. This is a limitation of Hibernate and there's also a ticket for it here, so the workaround is to use a different thread for the default DB query and force Hibernate to recreate the connection with the desired tenant. This workaround (which is taken from here) is necessary only for the login part. There aren't many reasons to switch tenants in the middle of the flow.

In this example, I created a Callable that's called TenantResolver, which contains the logic of querying the default schema in order to get the user's tenantId. 

@RequestMapping(value = "login", method = RequestMethod.POST)
    public ResponseEntity<?> createAuthenticationToken(@RequestBody JwtAuthenticationRequest authenticationRequest) throws AuthenticationException {
//Resolve the user's tenantId
        try {
            tenantResolver.setUsername(authenticationRequest.getUsername());
            ExecutorService es = Executors.newSingleThreadExecutor();
            Future<UserTenantRelation> utrFuture = es.submit(tenantResolver);
            UserTenantRelation utr = utrFuture.get();
            //TODO: handle utr == null, user is not found
            //Got the tenant, now switch to the context
            TenantContext.setCurrentTenant(utr.getTenant());
        } catch (Exception e) {
            e.printStackTrace();
        }

        // Perform the authentication
        final Authentication authentication = authenticationManager.authenticate(
                new UsernamePasswordAuthenticationToken(
                        authenticationRequest.getUsername(),
                        authenticationRequest.getPassword()
                )
        );
        SecurityContextHolder.getContext().setAuthentication(authentication);

        // Reload password post-security so we can generate token
        final User user = (User)userDetailsService.loadUserByUsername(authenticationRequest.getUsername());
        final String token = jwtTokenUtil.generateToken(user);

        // Return the token
        return ResponseEntity.ok(new JwtAuthenticationResponse(token, user));
    }

I didn't attach all the code and didn't create a dedicated GitHub repo for it. If it's requested, I will!