Hibernate Schema-based Multi-Tenancy using StatementInspector

blog-post-img
By Simon Martinelli
08.01.2024

When using schema-based multi-tenancy with Hibernate, a simple way to define the schema is using a MultiTenantConnectionProvider where we call SET SCHEMA on the connection. However, this approach does not work with certain databases like MS SQL.

@Override
public Connection getConnection(String tenantIdentifier) throws SQLException {
    final Connection connection = getAnyConnection();
    connection.createStatement()
            .execute(String.format("SET SCHEMA \"%s\";", tenantIdentifier));
    return connection;
}

Implementing StatementInspector for Schema-Based Multi-Tenancy

In this project, Spring Boot and Hibernate are used, so we were looking for a solution for those frameworks.

Hibernate provides the possibility to create a StatementInterceptor to manipulate the SQL statement and replace the DEFAULT_SCHEMA with the actual tenant schema. 

public class MultiTenantSchemaStatementInspector implements StatementInspector {

  private static final String SCHEMA_PART = "%s.";

  @Override
  public String inspect(String sql) {
    return sql.replace(SCHEMA_PART.formatted(DEFAULT_SCHEMA), 
                       SCHEMA_PART.formatted(TenantContext.getCurrentTenantSchema()));
  }
}

The StatementInterceptor must be registered with Hibernate. In Spring Boot this can be done by setting a property. The class name must be fully qualified. c.e. is the package name in our case.

spring.jpa.properties.hibernate.session_factory.statement_inspector=c.e.MultiTenantSchemaStatementInspector

TenantContext for Managing Tenant Schema

In the StatementInterceptor a class called TenantContext is used. This class used ThreadLocal to store the current tenant schema: 

public class TenantContext {

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

  private TenantContext() {
  }

  public static String getCurrentTenantSchema() {
    return currentTenantSchema.get();
  }

  public static void setCurrentTenantSchema(String tenant) {
    currentTenantSchema.set(tenant);
  }

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

Pre- and Post-Handling with HandlerInterceptor

Next, we must set the schema name to the TenantContext. For this, we’ve implemented a Spring HandlerInterceptor. What’s the difference between a HandlerInterceptor and a Filter? Let’s check the JavaDoc:

HandlerInterceptor is basically similar to a Servlet Filter, but in contrast to the latter it just allows custom pre-processing with the option of prohibiting the execution of the handler itself, and custom post-processing. Filters are more powerful, for example they allow for exchanging the request and response objects that are handed down the chain.

We use the preHandle method to set the schema name and the postHandle to clear it:

@Component
public class MultiTenantRequestInterceptor implements HandlerInterceptor {

  @Override
  public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
    var schmeaName = request.getHeader("TENANT-SCHEMA");

    TenantContext.setCurrentTenantSchema(schemaName;

    return true;
  }

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

Default Schema Definition and Considerations

As you may noticed in the code of the SchemaInspector implementation, the DEFAULT_SCHEMA is replaced. But where does the DEFAULT_SCHEMA come from? We must define it on the entities wherever it must be used, like tables, sequences, or join tables:

@Entity
@Table(schema = EntityMetadataConstants.DEFAULT_SCHEMA, name = "o_billing_run")
public class BillingRunEntity {

  @Id
  @SequenceGenerator(schema = EntityMetadataConstants.DEFAULT_SCHEMA, name = "seq_o_billing_run", sequenceName = "seq_o_billing_run")
  @GeneratedValue(strategy = GenerationType.SEQUENCE, generator = "seq_o_billing_run")
  private Long id;

  ...
}

Conclusion

This solution works, but it’s not perfect and is error-prone. You have to make sure that the default schema is defined in all entities. When using native SQL you also must include the DEFAULT_SCHEMA. But finally, it gets the job done and runs without problems in production for already a while.

Exploring the Potential of Schema-Based Multi-Tenancy for Your Project?

Intrigued by the possibility of implementing this approach in your Spring Boot project, or simply curious to delve deeper? Contact me to explore the potential of schema-based multi-tenancy for your project.

Related Posts

Simon Martinelli
Follow Me
Latest posts by Simon Martinelli (see all)