Tenant Per Schema with Spring Boot

In a recent small project, I'm using a Postgres database and I needed to create a reporting API where the data was stored in a "store per schema" situation. While I couldn't find "official" Spring documentation for doing multi-tenant architecture I did find quite a bit of documentation on "tenant per database" examples but very little "tenant per schema" information. After a lot of googling, I've pieced together something that works quite well.

Working example can be found on my github account.

Spring has built in support for multi-tenant design so no 3rd party libraries are required. There are 3 basic settings you'll need to configure:

  • hibernate.multiTenancy
  • hibernate.multitenantconnection_provider
  • hibernate.tenantidentifierresolver

While these settings can be configured in application properties, I found out that this won't work because of the order in which the EntityManager is wired up.

The first thing we'll need is to configure the EntityManager via code which is actually pretty straightforward. I created a class called HibernateConfig which contains the following:

@Configuration
public class HibernateConfig {

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

    @Bean
    public LocalContainerEntityManagerFactoryBean entityManagerFactory(DataSource dataSource, MultiTenantConnectionProvider multiTenantConnectionProvider,
                                                                       CurrentTenantIdentifierResolver tenantIdentifierResolver) {
        LocalContainerEntityManagerFactoryBean em = new LocalContainerEntityManagerFactoryBean();
        em.setDataSource(dataSource);
        em.setPackagesToScan("com.gdb.entity");

        em.setJpaVendorAdapter(this.jpaVendorAdapter());

        Map<String, Object> jpaProperties = new HashMap<>();
        jpaProperties.put(Environment.MULTI_TENANT, MultiTenancyStrategy.SCHEMA);
        jpaProperties.put(Environment.MULTI_TENANT_CONNECTION_PROVIDER, multiTenantConnectionProvider);
        jpaProperties.put(Environment.MULTI_TENANT_IDENTIFIER_RESOLVER, tenantIdentifierResolver);
        jpaProperties.put(Environment.FORMAT_SQL, true);

        em.setJpaPropertyMap(jpaProperties);
        return em;
    }
}

A couple of things to note:

  • I set the hibernate.multiTenancy property via Environment.MULTI_TENANT to SCHEMA
  • The MultiTenantConnectionProvider and CurrentTenantIdentifierResolver are injected because we are going to create custom versions of these classes as Components.

Another class I had to create was a subclass of MultiTenantConnectionProvider which looks like this:

@Component
public class SchemaPerTenantConnectionProvider implements MultiTenantConnectionProvider {

    @Autowired
    private DataSource dataSource;

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

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

    @Override
    public Connection getConnection(String tenantIdentifier) throws SQLException {
        final Connection connection = this.getAnyConnection();
        try {
            connection.createStatement().execute("SET search_path to " + tenantIdentifier);
        } catch (SQLException e) {
            throw new HibernateException("Could not alter JDBC connection to specified schema [" + tenantIdentifier + "]",
                    e);
        }
        return connection;
    }

    @Override
    public void releaseConnection(String tenantIdentifier, Connection connection) throws SQLException {
        try {
            connection.createStatement().execute("SET search_path to public");
        } catch (SQLException e) {
            throw new HibernateException("Could not alter JDBC connection to specified schema [" + tenantIdentifier + "]",
                    e);
        }

        connection.close();
    }

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

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

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

This class is passed a tenantIdentifier when a connection is requested and then I'm able to make sure that connection is going to use the correct schema. Since I'm using Postgres I have to modify the search_path. When the connection is released, I reset it back to the public schema since, in our case, the public schema is empty so a simple exception would be thrown with no risk of accessing insecure data. Because SchemaPerTenantConnectionProvider is a Component, Spring injects this into our LocalContainerEntityManagerFactoryBean, like I said before.

So how does the tenantIdentifier get passed to the provider? That's done with out sublcass of CurrentTenantIdentifierResolver which looks like the following:

@Component
public class HeaderTenantIdentifierResolver implements CurrentTenantIdentifierResolver {

    @Value("${multitenant.tenantKey}")
    String tenantKey;

    @Value("${multitenant.defaultTenant}")
    String defaultTenant;

    @Override
    public String resolveCurrentTenantIdentifier() {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        if (requestAttributes != null) {
            String tenantId = (String) requestAttributes.getAttribute(tenantKey, RequestAttributes.SCOPE_REQUEST);
            if (tenantId != null) {
                return tenantId;
            }
        }
        return defaultTenant;
    }

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

In our case, I'm putting the requested tenant in the Header of all the API requests and making sure that gets put in the RequestAttributes using a Filter (which I'll show next). If no tenant was provided, I just pass back the default which is public. Once again, this gets injected into the LocalContainerEntityManagerFactoryBean for us.

As I said, I created a Filter to take the Header out of the request and put it in RequestAttributes. Here's what that class looks like:

public class MultiTenantFilter implements Filter {

    @Value("${multitenant.tenantKey}")
    String tenantKey;

    @Value("${multitenant.defaultTenant}")
    String defaultTenant;

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {

    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest req = (HttpServletRequest) request;
        String tenant = req.getHeader(tenantKey);

        if (tenant != null) {
            req.setAttribute(tenantKey, tenant);
        } else {
            req.setAttribute(tenantKey, defaultTenant);
        }
        chain.doFilter(request, response);
    }

    @Override
    public void destroy() {

    }
}

And then to wire that up into the application, I did the following in my Application class:

@SpringBootApplication
public class Application {

    @Autowired
    AutowireCapableBeanFactory beanFactory;

    public static void main(String[] args) {
        SpringApplication.run(Application.class, args);
    }

    @Bean
    public FilterRegistrationBean myFilter() {
        FilterRegistrationBean registration = new FilterRegistrationBean();
        Filter tenantFilter = new MultiTenantFilter();
        beanFactory.autowireBean(tenantFilter);
        registration.setFilter(tenantFilter);
        registration.addUrlPatterns("/*");
        return registration;
    }
}

For completeness, here's my application.yml file:

multitenant:  
  tenantKey: X-Person-Tenant
  defaultTenant: public
spring:  
  datasource:
      url: jdbc:postgresql://localhost:5432/multi_tenant
      username: ****
      password: ****
      driverClassName: org.postgresql.Driver
  jpa:
    database-platform: org.hibernate.dialect.PostgreSQLDialect
    generate-ddl: false

I ran the following INSERT statements on my database:

insert into schema1.person (first_name, last_name) values ('Gregg', 'Bolinger');  
insert into schema2.person (first_name, last_name) values ('John', 'Smith');  

Using schema1 in the Header, I get the following

Using schema2 in the Header, I get the following:

If you see anything that I might have done that can be improved on, please let me know.