How to write a custom Spring Data Base Repository

Imagine having a tool that can automatically detect JPA and Hibernate performance issues. Wouldn’t that be just awesome?

Well, Hypersistence Optimizer is that tool! And it works with Spring Boot, Spring Framework, Jakarta EE, Java EE, Quarkus, or Play Framework.

So, enjoy spending your time on the things you love rather than fixing performance issues in your production system on a Saturday night!

Introduction

In this article, we are going to see how we can write a custom Spring Data base Repository that you could use instead of the default ones, like the overly common JpaRepository.

Now, why would you even want to do that? Well, most of the default repositories extend the CrudRepository, which provides some very questionable defaults, like findAll or deleteAll that shouldn’t be really inherited by every single data access Repository instance.

Spring Data Repository

I’ve been a long-time Spring user. The first time I used Spring was at the end of 2004.

There are two main reasons why I’ve been using Spring for such a long time:

  • their documentation is amazing
  • both the Core and the other modules are highly customizable

For example, not only that the Spring Data Repository abstraction is very well-documented, but the development team made it possible for us to customize it.

In the long run, it’s the ability to customize a given framework that will provide you with the best return on investment. Having lots of features are nice, but being able to change the default behavior is great.

Meet the BaseJpaRepository and BaseHibernateRepository

As I explained in this article, I don’t like to have all my JPA Repository instances inherit the findAll or deleteAll methods because I believe these methods could be miss used.

So, instead of using the default JpaRepository, I’d rather use the following BaseJpaRepository and BaseHibernateRepository alternatives:

Custom Spring Data base Repository

So, instead of inheriting from the default JpaRepository, my data access objects extend the BaseJpaRepository and BaseHibernateRepository:

@Repository
public interface PostRepository extends BaseJpaRepository<Post, Long>,
    BaseHibernateRepository<Post> {
}

The BaseJpaRepository looks as follows:

@NoRepositoryBean
public interface BaseJpaRepository<T, ID> extends Repository<T, ID>, 
    QueryByExampleExecutor<T> {

    Optional<T> findById(ID id);

    boolean existsById(ID id);

    T getReferenceById(ID id);

    List<T> findAllById(Iterable<ID> ids);

    long count();

    void delete(T entity);

    void deleteAllInBatch(Iterable<T> entities);

    void deleteById(ID id);

    void deleteAllByIdInBatch(Iterable<ID> ids);

    void flush();
}

First, we extend the Spring Data Repository interface so that the framework can generate the actual implementation for the data access methods we will declare explicitly.

We can also extend the QueryByExampleExecutor interface as the findAll methods provided by this interface allow users to provide the mandatory filtering criteria.

Next, we define which methods we want to be inherited by all our data access Repository specific classes.

The only findAll and deleteAll methods we included are the ones that provide a way to limit the number of entities we are about to fetch or remove.

However, there’s no save method in the BaseJpaRepository because, as I explained in this article, there’s no such thing in the JPA specification. On the other hand, we have persist and merge, and the Hibernate-specific update.

So, for this reason, our Repository abstractions extend the BaseHibernateRepository interface, which looks like this:

public interface BaseHibernateRepository<T> {

    /*
        The persist methods are meant to save the newly created entities.
    */
    
    <S extends T> S persist(S entity);

    <S extends T> S persistAndFlush(S entity);

    <S extends T> List<S> persistAll(Iterable<S> entities);

    <S extends T> List<S> peristAllAndFlush(Iterable<S> entities);

    /*
        The merge methods are meant to propagate detached entity state changes
        if they are really needed.
    */

    <S extends T> S merge(S entity);

    <S extends T> S mergeAndFlush(S entity);

    <S extends T> List<S> mergeAll(Iterable<S> entities);

    <S extends T> List<S> mergeAllAndFlush(Iterable<S> entities);

    /*
        The update methods are meant to force the synchronization of the 
        detached entity state changes.
    */

    <S extends T> S update(S entity);

    <S extends T> S updateAndFlush(S entity);

    <S extends T> List<S> updateAll(Iterable<S> entities);

    <S extends T> List<S> updateAllAndFlush(Iterable<S> entities);
}

While Spring Data JPA can provide an implementation for the methods declared by the BaseJpaRepository interface, it doesn’t know how to do that for the ones we defined by the BaseHibernateRepository interface.

And as explained by this article, this can be handled by the custom Repository feature.

Therefore, we can provide the following BaseHibernateRepositoryImpl implementation for this interface:

public class BaseHibernateRepositoryImpl<T> implements BaseHibernateRepository<T> {

    @PersistenceContext
    private EntityManager entityManager;

    public <S extends T> S persist(S entity) {
        entityManager.persist(entity);
        return entity;
    }

    public <S extends T> S persistAndFlush(S entity) {
        persist(entity);
        entityManager.flush();
        return entity;
    }

    public <S extends T> List<S> persistAll(Iterable<S> entities) {
        List<S> result = new ArrayList<>();
        for(S entity : entities) {
            result.add(persist(entity));
        }
        return result;
    }

    public <S extends T> List<S> peristAllAndFlush(Iterable<S> entities) {
        return executeBatch(() -> {
            List<S> result = new ArrayList<>();
            for(S entity : entities) {
                result.add(persist(entity));
            }
            entityManager.flush();
            return result;
        });
    }

    public <S extends T> S merge(S entity) {
        return entityManager.merge(entity);
    }

    public <S extends T> S mergeAndFlush(S entity) {
        S result = merge(entity);
        entityManager.flush();
        return result;
    }

    public <S extends T> List<S> mergeAll(Iterable<S> entities) {
        List<S> result = new ArrayList<>();
        for(S entity : entities) {
            result.add(merge(entity));
        }
        return result;
    }

    public <S extends T> List<S> mergeAllAndFlush(Iterable<S> entities) {
        return executeBatch(() -> {
            List<S> result = new ArrayList<>();
            for(S entity : entities) {
                result.add(merge(entity));
            }
            entityManager.flush();
            return result;
        });
    }

    public <S extends T> S update(S entity) {
        session().update(entity);
        return entity;
    }

    public <S extends T> S updateAndFlush(S entity) {
        update(entity);
        entityManager.flush();
        return entity;
    }

    public <S extends T> List<S> updateAll(Iterable<S> entities) {
        List<S> result = new ArrayList<>();
        for(S entity : entities) {
            result.add(update(entity));
        }
        return result;
    }

    public <S extends T> List<S> updateAllAndFlush(Iterable<S> entities) {
        return executeBatch(() -> {
            List<S> result = new ArrayList<>();
            for(S entity : entities) {
                result.add(update(entity));
            }
            entityManager.flush();
            return result;
        });
    }

    protected Integer getBatchSize(Session session) {
        SessionFactoryImplementor sessionFactory = session
            .getSessionFactory()
            .unwrap(SessionFactoryImplementor.class);
            
        final JdbcServices jdbcServices = sessionFactory
            .getServiceRegistry()
            .getService(JdbcServices.class);
            
        if(!jdbcServices.getExtractedMetaDataSupport().supportsBatchUpdates()) {
            return Integer.MIN_VALUE;
        }
        return session
            .unwrap(AbstractSharedSessionContract.class)
            .getConfiguredJdbcBatchSize();
    }

    protected <R> R executeBatch(Supplier<R> callback) {
        Session session = session();
        Integer jdbcBatchSize = getBatchSize(session);
        Integer originalSessionBatchSize = session.getJdbcBatchSize();
        try {
            if (jdbcBatchSize == null) {
                session.setJdbcBatchSize(10);
            }
            return callback.get();
        } finally {
            session.setJdbcBatchSize(originalSessionBatchSize);
        }
    }

    protected Session session() {
        return entityManager.unwrap(Session.class);
    }
}

Testing Time

First, we will create a ForumService interface that looks like this:

public interface ForumService {

    Post findById(Long id);

    Post createPost(Post post);

    Post updatePost(Post post);
}

And we provide the following ForumServiceImpl implementation for it:

@Service
@Transactional(readOnly = true)
public class ForumServiceImpl implements ForumService {

    private PostRepository postRepository;

    public ForumServiceImpl(
            @Autowired PostRepository postRepository) {
        this.postRepository = postRepository;
    }

    public Post findById(Long id) {
        return postRepository.findById(id).orElse(null);
    }

    @Transactional
    @Override
    public Post createPost(Post post) {
        return postRepository.persist(post);
    }

    @Transactional
    @Override
    public Post updatePost(Post post) {
        postRepository.update(post);
        return post;
    }
}

Notice that we define the @Transactional(readOnly = true) at the class level, so all methods are read-only by default.

However, for read-write methods, we explicitly declare them with the @Transactional annotation. For more details about this pattern, check out this article.

With the ForumService in place, if we run the following test case:

Long postId = forumService.createPost(
    new Post()
        .setId(1L)
        .setTitle("High-Performance Java Persistence")
).getId();

Post post = forumService.findById(postId);
assertEquals("High-Performance Java Persistence", post.getTitle());

post.setTitle("High-Performance Java Persistence, 2nd edition");
forumService.updatePost(post);

Hibernate will execute the following SQL statements:

INSERT INTO post (
    title, 
    id
) 
VALUES (
    'High-Performance Java Persistence', 
    1
)

SELECT 
    p.id as id1_1_0_, 
    p.title as title2_1_0_ 
FROM 
    post p 
WHERE 
    p.id= 1

UPDATE 
    post 
SET 
    title = 'High-Performance Java Persistence, 2nd edition'
WHERE 
    id = 1

Cool, right?

If you enjoyed this article, I bet you are going to love my Book and Video Courses as well.

Seize the deal! 40% discount. Seize the deal! 40% discount.

Conclusion

Writing a custom Spring Data base Repository is actually very easy.

The advantage of defining your own Spring Data base Repository is that you can remove all the base methods that you will not need.

Transactions and Concurrency Control eBook

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.