Montag, 7. Dezember 2015

Multi-tenant cloud applications with Spring-Data-JPA and EclipseLink

We have been working on an application that requires a multi-tenancy architecture, tenants are created on the fly, potentially reaching a couple hundred. This mean, that multi-tenancy via database/schema separation is off the table, which leaves us with a single table approach and discriminator columns.

Usually Hibernate is our tool of choice for JPA persistence and Hibernate4 added multi-tenancy support - great! After fiddling around for a while and reading the documentation a second time, you will find a comment for the discriminator strategy: "This strategy is not yet implemented in Hibernate as of 4.0 and 4.1. Its support is planned for 5.0.".

Too bad, what now? Well since JPA 2.1 adds multi-tenancy support and EclipseLink is the reference implementation, maybe they support our need? Turns out they do since version 2.3.0 and they have a great documentation which has more than enough information to get you started.

The last piece of the puzzle is Spring-Data-JPA for the data-access-layer, one of our favorites - and it works with EclipseLink out of the box. Now we just have to make it behave nicely with multi-tenant entities.

For multi-tenancy to work EclipseLink requires a property on the PersistenceUnit/EntityManager to be set. This can be done by calling
em.setProperty(PersistenceUnitProperties.MULTITENANT_PROPERTY_DEFAULT, currentTenantId);
at runtime. With Spring-Data you don't need to inject the EntityManager anymore, you don't even need to know there is one being used under the hood, so where can this be done?

First things first

To get started we define an interface that we can implement to access the id of the current tenant:
public interface CurrentTenantResolver<T extends Serializable> {
    T getCurrentTenantId();
}
In our case we use Apache Shiro as security framework and access the current tenant from the session.
public class ShiroCurrentTenantResolver
        implements CurrentTenantResolver<Long> {
    @Override
    public Long getCurrentTenantId() {
        Session session = SecurityUtils.getSubject().getSession();
        return (Long) session.getAttribute("tenantId");
    }
}

Approach 1: Custom EntityMangerFactory

There was a blog post, which I can't seem to find anymore at this time, that suggested to use a custom EntityManagerFactory and Spring-EntityManagerFactoryBean to set the tenant-id property like so:
public class TenantAwareEntityManagerFactory implements EntityManagerFactory {
    private final EntityManagerFactory delegate;
    private final CurrentTenantResolver<Long> resolver;

    public TenantAwareEntityManagerFactory(EntityManagerFactory delegate,
                                           CurrentTenantResolver<Long> resolver) {
        this.delegate = delegate;
        this.resolver = resolver;
    }
    
    @Override
    public EntityManager createEntityManager() {
        Long tenantID = resolver.getCurrentTenantId();
        Map<String, Long> map = new HashMap<String, Long>();
        map.put(PersistenceUnitProperties.MULTITENANT_PROPERTY_DEFAULT, tenantID);
        return delegate.createEntityManager(map);
    }
    
    @Override
    public EntityManager createEntityManager(Map map) {
        Long tenantID = resolver.getCurrentTenantId();
        map.put(PersistenceUnitProperties.MULTITENANT_PROPERTY_DEFAULT, tenantID);
        return delegate.createEntityManager(map);
    }
    // delegate the rest
}
This would have been a really nice and clean solution but unfortunately this doesn't work in all cases - specifically it doesn't work when you use transactions via Spring's @Transactional annotation. Why is that?

EclipseLink requires the tenant property to be set on the EntityManager after the transaction begins. This can be seen in all the examples in the documentation but is not clearly stated. There are however a couple entries in the mailing-list like this one that point to the problem.

When using @Transactional Spring handles the call to EntityManager.beginTransaction() and in order to do that it has to create an EntityManager instance thus setting the property before the transaction begins.

Approach 2: Custom Repository/RepositoryFactory implementation

In order to get this to work we have to set the tenant property inside the repository and we have to do it for every call to a repository method. You can probably achieve this with some AOP magic but here is the pragmatic way - extend SimpleJpaRepository and override all query methods:
public class MultiTenantSimpleJpaRepository<T, ID extends Serializable> extends SimpleJpaRepository<T, ID> {
    private final CurrentTenantResolver tenantResolver;
    private final EntityManager em;

    public MultiTenantSimpleJpaRepository(JpaEntityInformation<T, ?> entityInformation,
                                          EntityManager em, CurrentTenantResolver tenantResolver) {
        super(entityInformation, em);
        this.tenantResolver = tenantResolver;
        this.em = em;
    }

    public MultiTenantSimpleJpaRepository(Class<T> domainClass, EntityManager em,
                                          CurrentTenantResolver tenantResolver) {
        super(domainClass, em);
        this.tenantResolver = tenantResolver;
        this.em = em;
    }

    protected void setCurrentTenant() {
        em.setProperty(PersistenceUnitProperties.MULTITENANT_PROPERTY_DEFAULT, tenantResolver.getCurrentTenantId());
    }

    @Override
    public <S extends T> S save(S entity) {
        setCurrentTenant();
        return super.save(entity);
    }
    // override the other methods
}
If you use the awesome QueryDslRepository then you should subclass it as well. To get Spring-Data-JPA to use our custom repository implementations we need to create our own RepositoryFactory and RepositoryFactoryBean.
public class MultiTenantJpaRepositoryFactory extends JpaRepositoryFactory {
    private final CurrentTenantResolver currentTenantResolver;

    public MultiTenantJpaRepositoryFactory(EntityManager entityManager, CurrentTenantResolver currentTenantResolver) {
        super(entityManager);
        this.currentTenantResolver = currentTenantResolver;
    }

    @Override
    @SuppressWarnings("unchecked")
    protected JpaRepository<?, ?> getTargetRepository(RepositoryMetadata metadata, EntityManager entityManager) {
        final Class repositoryInterface = metadata.getRepositoryInterface();
        final JpaEntityInformation<?, Serializable> entityInformation = getEntityInformation(metadata.getDomainType());
        
        final SimpleJpaRepository<?, ?> repo = isQueryDslExecutor(repositoryInterface) ?
                new MultiTenantQueryDslJpaRepository(entityInformation, entityManager, currentTenantResolver) :
                new MultiTenantSimpleJpaRepository(entityInformation, entityManager, currentTenantResolver);
        repo.setLockMetadataProvider(LockModeRepositoryPostProcessor.INSTANCE.getLockMetadataProvider());
        return repo;
    }

    @Override
    protected Class<?> getRepositoryBaseClass(RepositoryMetadata metadata) {
        if (isQueryDslExecutor(metadata.getRepositoryInterface())) {
            return MultiTenantQueryDslJpaRepository.class;
        } else {
            return MultiTenantSimpleJpaRepository.class;
        }
    }

    private boolean isQueryDslExecutor(Class<?> repositoryInterface) {
        return QUERY_DSL_PRESENT && QueryDslPredicateExecutor.class.isAssignableFrom(repositoryInterface);
    }
}
public class MultiTenantJpaRepositoryFactoryBean<T extends Repository<S, ID>, S, ID extends Serializable> extends JpaRepositoryFactoryBean<T, S, ID> {
    private CurrentTenantResolver currentTenantResolver;

    @Override
    protected RepositoryFactorySupport createRepositoryFactory(EntityManager entityManager) {
        return new MultiTenantJpaRepositoryFactory(entityManager, currentTenantResolver);
    }

    @Override
    public void afterPropertiesSet() {
        Assert.notNull(currentTenantResolver, "CurrentTenantResolver must not be null!");
        super.afterPropertiesSet();
    }

    @Autowired
    public void setCurrentTenantResolver(CurrentTenantResolver currentTenantResolver) {
        this.currentTenantResolver = currentTenantResolver;
    }
}
And finally we have to tell Spring-Data-JPA to use our new factory by specifying the factory class in your Spring XML:
<jpa:repositories base-package="com.codecraft.server.repository"
    factory-class="com.codecraft.server.orm.MultiTenantJpaRepositoryFactoryBean"/>
Now we can use all the features of Spring-Data-JPA in a multi-tenancy environment, EclipseLink will automatically modify all queries except JPA native queries to return only results for the current tenant.
Also EclipseLink is new to us, it has been really easy and fun to use so far. It is gonna be a tough choice between Hibernate and EclipseLink in the future.

Edit:

As it turns out, this is much simpler if you extend Springs's JpaTransactionManager and set the property there. No need for the MultiTenantSimpleJpaRepository/MultiTenantSimpleJpaRepositoryFactory anymore.