Jakarta EE: Multitenancy with JPA on WildFly, Part 2

This is the second half of a two-part article about multitenancy with the Jakarta Persistence API (JPA) on WildFly. In Part 1, I showed you how to implement multitenancy using a database. In Part 2, I'll show you how to implement multitenancy using a schema and the Jakarta Persistence API (JPA) on WildFly. You'll learn how to implement JPA's CurrentTenantIdentifierResolver and MultiTenantConnectionProvider interfaces, and how to use JPA’s persistence.xml file to configure the required classes based on these interfaces.

Implementation code

The first part of this article offers a conceptual overview of multitenancy with JPA on WildFly and an example of multitenancy using a database. This second part shifts the focus to multitenancy using a schema and JPA. In this case, I assume that WildFly manages the data source and connection pool and that EJB (Jakarta Enterprise Beans) handles the container-managed transactions.

Two interfaces for multitenancy

As I explained in my previous article, two interfaces are crucial for implementing multitenancy in JPA and Hibernate:

• The MultiTenantConnectionProvider interface is responsible for connecting tenants to their respective databases and services. We will use this interface and a tenant identifier to switch between databases for different tenants.
• CurrentTenantIdentifierResolver is responsible for identifying the tenant. We will use this interface to define what is considered a tenant (more about this later) and to provide the correct tenant identifier to MultiTenantConnectionProvider.

Next, we'll look at the three classes that implement these interfaces.

SchemaMultiTenantProvider

SchemaMultiTenantProvider is an implementation of the MultiTenantConnectionProvider interface. This class contains logic to switch to the schema that matches the given tenant identifier.

The SchemaMultiTenantProvider class also implements the ServiceRegistryAwareService, which allows us to inject a service during the configuration phase. Here’s the code for the SchemaMultiTenantProvider class:

public class SchemaMultiTenantProvider implements MultiTenantConnectionProvider, ServiceRegistryAwareService {

    private static final long serialVersionUID = 1L;
    private static final String TENANT_SUPPORTED = "SCHEMA";
    private DataSource dataSource;
    private String typeTenancy ;

    @Override
    public boolean supportsAggressiveRelease() {
        return false;
    }
    @Override
    public void injectServices(ServiceRegistryImplementor serviceRegistry) {
        typeTenancy = (String) ((ConfigurationService)serviceRegistry
                .getService(ConfigurationService.class))
                .getSettings().get("hibernate.multiTenancy");

        dataSource = (DataSource) ((ConfigurationService)serviceRegistry
                .getService(ConfigurationService.class))
                .getSettings().get("hibernate.connection.datasource");
    }

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

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

    @Override
    public Connection getAnyConnection() throws SQLException {
        final Connection connection = dataSource.getConnection();
        resetConnection(connection);// To make sure the connection start using schema/database default.
        return connection;
    }

    @Override
    public Connection getConnection(String tenantIdentifier) throws SQLException {
        //Just use the multitenancy if the hibernate.multiTenancy == SCHEMA
        if(TENANT_SUPPORTED.equals(typeTenancy)) {
            try {
                final Connection connection = getAnyConnection();
                connection.createStatement().execute("SET SCHEMA '" + tenantIdentifier + "'");
                return connection;
            } catch (final SQLException e) {
                 throw new HibernateException("Error trying to alter schema [" + tenantIdentifier + "]", e);
            }
        }

        return getAnyConnection();
    }



    @Override
    public void releaseAnyConnection(Connection connection) throws SQLException {
        //As the Wildfly/JBoss has the Container-Managed Container change the SCHEMA in the end can be dangerous (SET SCHEMA 'public').
        //Thus it just closes the connection.
        connection.close();
    }

    private void resetConnection(Connection connection){
        if(TENANT_SUPPORTED.equals(typeTenancy)) {
            try {
                connection.createStatement().execute("SET SCHEMA 'public'");
            } catch (final SQLException e) {
                throw new HibernateException("Error trying to alter schema [public]", e);
            }
        }
    }

    @Override
    public void releaseConnection(String tenantIdentifier, Connection connection) throws SQLException {
        releaseAnyConnection(connection);
    }

}

As you can see, we call the injectServices method to populate the datasource and typeTenancy attributes. We use the datasource attribute to get a connection from the data source, and we use the typeTenancy attribute to find out if the class supports the multiTenancy type. We call the getConnection method to get a data source connection. This method uses the tenant identifier to locate and switch to the correct schema.

MultiTenantResolver

MultiTenantResolver is a simple, abstract class that implements the CurrentTenantIdentifierResolver interface. This class aims to provide a setTenantIdentifier method to all CurrentTenantIdentifierResolver implementations:

public abstract class MultiTenantResolver implements CurrentTenantIdentifierResolver {

    protected String tenantIdentifier;

    public void setTenantIdentifier(String tenantIdentifier) {
        this.tenantIdentifier = tenantIdentifier;
    }
}

We only use this class to provide the setTenantIdentifier method.

SchemaTenantResolver

SchemaTenantResolver also implements the CurrentTenantIdentifierResolver interface. This class is the concrete class of MultiTenantResolver:

public class SchemaTenantResolver extends MuiltiTenantResolver {

    private Map<String, String> userDatasourceMap;

    public SchemaTenantResolver(){
        userDatasourceMap = new HashMap();
        userDatasourceMap.put("default", "public");
        userDatasourceMap.put("username1", "usernameone");
        userDatasourceMap.put("username2", "usernametwo");

    }

    @Override
    public String resolveCurrentTenantIdentifier() {
        if(this.tenantIdentifier != null
                && userDatasourceMap.containsKey(this.tenantIdentifier)){
            return userDatasourceMap.get(this.tenantIdentifier);
        }

        return userDatasourceMap.get("default");
    }

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

}

Notice that SchemaTenantResolver uses a Map to define the correct schema for a given tenant. The tenant, in this case, is mapped by the user.

Configure and define the tenant

Now, we need to use JPA's persistence.xml file to configure the tenant:

<persistence>
    <persistence-unit name="jakartaee8">

        <jta-data-source>jdbc/MyDataSource</jta-data-source>
        <properties>
            <property name="javax.persistence.schema-generation.database.action" value="none" />
            <property name="hibernate.dialect" value="org.hibernate.dialect.PostgresPlusDialect"/>
            <property name="hibernate.multiTenancy" value="SCHEMA"/>
            <property name="hibernate.tenant_identifier_resolver" value="net.rhuanrocha.dao.multitenancy.SchemaTenantResolver"/>
            <property name="hibernate.multi_tenant_connection_provider" value="net.rhuanrocha.dao.multitenancy.SchemaMultiTenantProvider"/>
        </properties>

    </persistence-unit>
</persistence>

We define the tenant in the JPA and Hibernate EntityManagerFactory:

@PersistenceUnit
protected EntityManagerFactory emf;


protected EntityManager getEntityManager(String multitenancyIdentifier){

    final MuiltiTenantResolver tenantResolver = (MuiltiTenantResolver) ((SessionFactoryImplementor) emf).getCurrentTenantIdentifierResolver();
    tenantResolver.setTenantIdentifier(multitenancyIdentifier);

    return emf.createEntityManager();
}

Note that we call the setTenantIdentifier before creating a new instance of EntityManager.

Conclusion

This article presented a simple example of multitenancy using a schema in a database. There are many ways to use a database for multitenancy. My point has been to show you how to implement the CurrentTenantIdentifierResolver and MultiTenantConnectionProvider interfaces. I've also shown you how to use JPA's persistence.xml to configure the required classes based on the two interfaces.

Keep in mind that for this example, I have assumed that WildFly manages the data source and connection pool and that we're using enterprise beans for the container-managed transactions. If you want to go deeper with this example, you can find the complete application code and further instructions on my GitHub repository.