Multitenancy breaks after 6.2 breaks upgrade

Hibernate ORM
1

TLDR please help set up schema-based multitenancy in Hibernate 6.2

Hello, I am upgrading a Spring Boot project from 2.7 to 3.1.5, which involves upgrading Hibernate 5 to 6.2. This has totally broken my schema-based multitenancy setup. I am using Postgres along with Spring Data JPA, with various entity classes and JpaRepository interfaces.

I have implementations of MultiTenantConnectionProvider and TenantIdentifierResolver registered as @Component Beans. In Hibernate 5, I was only able to get Hibernate multitenancy working by having each of these Beans implement HibernatePropertiesCustomizer in order for each class to register itself with Hibernate (registering in application properties had no effect). And I’m aware in the upgrade to 6 that the property MULTI_TENANT was removed).

The issue is that after the upgrade, when the connection provider registers itself, the Jdbc connection fails with “Connection is closed”. For example, if I run a test class using a Testcontainer that autowires a repository, it will error. Conversely, if I remove the registration method from the connection provider, I can get these repository tests to run in isolation, but then my project won’t build (with Maven) as the multitenancy setup will not be implemented.

This worked correctly before the upgrade, so I’m wondering what has gone wrong. I have tried registering the connection provider in other ways, such as a HibernateConfig Bean or in LocalContainerEntityManagerFactoryBean, but these just have the same effect as it registering itself with HibernatePropertiesCustomizer.

It seems to be a chicken-and-egg situation, but I’m looking for any insight, as presumably this kind of setup is possible. One more piece of the puzzle is that my connection provider autowires a custom @Configuration Bean which brings in application properties such as the schemata names etc.

Here is my connection provider (logs and names removed):

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.Optional;

import javax.sql.DataSource;

import org.[etc].config.TenantIdentificationConfig;
import org.[etc].exception.SchemaException;
import org.[etc].service.ThreadLocalSchemaManager;
import org.hibernate.cfg.AvailableSettings;
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.boot.autoconfigure.orm.jpa.HibernatePropertiesCustomizer;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;

@Component
public class TenantBasedConnectionProvider implements MultiTenantConnectionProvider, HibernatePropertiesCustomizer {

    private static final Logger LOGGER = LoggerFactory.getLogger(TenantBasedConnectionProvider.class);

    private final DataSource dataSource;
    private final TenantIdentificationConfig tenantIdentificationConfig;
    private final ThreadLocalSchemaManager threadLocalSchemaManager;

    @Autowired
    public TenantBasedConnectionProvider(final DataSource dataSource,
            final TenantIdentificationConfig tenantIdentificationConfig,
            final ThreadLocalSchemaManager threadLocalSchemaManager) {
        this.tenantIdentificationConfig = tenantIdentificationConfig;
        this.dataSource = dataSource;
        this.threadLocalSchemaManager = threadLocalSchemaManager;
    }

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

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

    @Override
    public Connection getConnection(final String tenantId) throws SQLException {
        final String threadLocalTenantId = threadLocalSchemaManager.getTenantId();
        if (!threadLocalSchemaManager.schemaNameHasValidTenant(tenantId)) {
            throw new SchemaException("Attempted to set db connection to invalid schema", tenantId);
        }

        try (final Connection connection = new JdbcTemplate(dataSource).getDataSource().getConnection()) {
            connection.setSchema(tenantId);
            return connection;
        }
    }

    @Override
    public void releaseConnection(final String tenantId, final Connection connection) throws SQLException {
        connection.setSchema(tenantIdentificationConfig.getDefaultSchema());
        connection.close();
    }

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

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

    @Override
    public <T> T unwrap(Class<T> aClass) {
        throw new UnsupportedOperationException("Can't unwrap this.");
    }

    @Override
    public void customize(Map<String, Object> hibernateProperties) {
        hibernateProperties.put(AvailableSettings.MULTI_TENANT_CONNECTION_PROVIDER, this);
    }
}

Here is my tenant identifier:

@Component
public class TenantIdentifierResolver implements CurrentTenantIdentifierResolver, HibernatePropertiesCustomizer {

    private static final Logger LOGGER = LoggerFactory.getLogger(TenantIdentifierResolver.class);

    private final ThreadLocalSchemaManager threadLocalSchemaManager;

    @Autowired
    public TenantIdentifierResolver(final ThreadLocalSchemaManager threadLocalSchemaManager) {
        this.threadLocalSchemaManager = threadLocalSchemaManager;
    }

    @Override
    public String resolveCurrentTenantIdentifier() {
        final String currentSchema = threadLocalSchemaManager.getTenantId();
        if (currentSchema != null) {
            return currentSchema;
        }

        throw new SchemaException("Could not resolve current schema");
    }

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

    @Override
    public void customize(Map<String, Object> hibernateProperties) {;
        hibernateProperties.put(AvailableSettings.MULTI_TENANT_IDENTIFIER_RESOLVER, this);
    }

}

Note that the ThreadLocalSchemaManager is a service that uses a ThreadLocal variable to store the tenantId.

The packages I am using are:

io.hypersistence.hypersistence-utils-hibernate-62
3.7.0

org.springframework.boot.spring-boot-starter-data-jpa
(gets its version from parent POM in Spring Boot 3.1.5)

Please let me know if there is a problem with my setup or if I have missed anything in my upgrade. I would like to see a functional example, but I only find examples with multiple data sources rather than multiple schemata in one database

2

I have solved it. It was a red herring. While debugging I had introduced the @Container annotation in my PostgresIntegrationTest interface that all Testcontainer tests implement. This annotation closes the Testcontainer after the first test runs. You cannot use @Container or @ServiceConnection - you have to use @DynamicPropertyRegistry and call container.start() within the method

Another issue was that in MultiTenantConnectionProvider.getConnection I had put the connection inside a try block’s parenthesis, which meant the connection would get closed automatically. You have to leave it open and allow Hibernate to manage it

1 Like