Count Queries In JPA Using CriteriaQuery

1. Introduction

Java Persistence API (JPA) is a widely used specification for accessing, persisting, and managing data between Java objects and a relational database. One common task in JPA applications is counting the number of entities that meet certain criteria. This task can be efficiently accomplished using the CriteriaQuery API provided by JPA.

The core components of Criteria Query are the CriteriaBuilder and CriteriaQuery interfaces. The CriteriaBuilder is a factory for creating various query elements such as predicates, expressions, and criteria queries. On the other hand, the CriteriaQuery represents a query object that encapsulates the selection, filtering, and ordering criteria.

In this article, we’ll look into count queries in JPA, exploring how to leverage the Criteria Query API to perform count operations with ease and efficiency. We’ll start with a quick overview of CriteriaQuery and how they can be used to produce count queries. Next, we’ll take a simple example of a library management system where we can leverage CriteriaQuery API to produce counts for the books in various scenarios.

2. Dependencies and Example Setup

Firstly, let’s ensure we have the required maven dependencies including spring-data-jpa, spring-boot-starter-test, and h2 as an in-memory database:

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-jpa</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-test</artifactId>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>com.h2database</groupId>
    <artifactId>h2</artifactId>
    <scope>runtime</scope>
</dependency>

Now that our dependencies are set, let’s introduce a simple example of a library management system. It will allow us to perform various queries such as counting all books counting books by a certain author, title, and year in various combinations. Let’s introduce a Book entity with the title, author, category, and year fields:

@Entity
public class Book {
    @Id
    @GeneratedValue(strategy = GenerationType.IDENTITY)
    private Long id;
    private String title;
    private String Category;
    private String author;
    private int year;

    // Constructors, getters, and setters
}

Next, let’s write a repository interface that will allow us to perform various operations with the Book entity:

public interface BookRepositoryCustom {
    long countAllBooks();
    long countBooksByTitle(String title);
    long countBooksByAuthor(String author);
    long countBooksByCategory(String category);
    long countBooksByTitleAndAuthor(String title, String author);
    long countBooksByAuthorOrYear(String author, int year);
}

3. Counting Entities With CriteriaQuery

Count queries are commonly used to determine the total number of entities that satisfy certain conditions. Using CriteriaQuery, we can construct count queries straightforwardly and efficiently. Let’s dive into a step-by-step guide on how to perform count queries using Criteria Query in JPA.

3.1. Initialize CriteriaBuilder and CriteriaQuery

To begin constructing our count query, we first need to obtain an instance of the CriteriaBuilder from the EntityManager. The CriteriaBuilder serves as our entry point for creating query elements:

CriteriaBuilder cb = entityManager.getCriteriaBuilder();

This object is used to construct different parts of the query, such as the criteria query, expressions, predicates, and selections. Let’s create a criteria query from it next:

CriteriaQuery<Long> cq = cb.createQuery(Long.class);

Here, we specify the result type of our query as Long, indicating that we expect the query to return a count value. 

3.2. Create Root and Select Count

Next, let’s create a Root object representing the entity for which we want to perform the count operation. We’ll then use the CriteriaBuilder to construct a count expression based on this Root object:

Root<Book> bookRoot = cq.from(Book.class);
cq.select(cb.count(bookRoot));

Here, first, we define the query’s root, specifying that the query is based on the Book entity. Next, we create a count expression using cb.count() provided by CriteriaBuilder. The count method counts the number of rows in the result of the query. It takes an expression (in this case, bookRoot) as an argument and returns an expression that represents the count of the number of rows that match the criteria defined in the query

Finally, cq.select() then sets the result of the query to be this count expression. Essentially, it tells the query that the final result should be the number of Book entities that match the specified conditions (if any which we will see later).

3.3. Execute the Query

Once the CriteriaQuery is constructed, we can execute the query using the EntityManager:

Long count = entityManager.createQuery(cq).getSingleResult();

Here, we use entityManager.createQuery(cq) to create a TypedQuery from the CriteriaQuery, and getSingleResult() to retrieve the count value as a single result.

4. Handling Criteria and Conditions

In real-world scenarios, count queries often involve filtering based on certain criteria or conditions. Criteria Query provides a flexible mechanism for adding criteria to our queries using predicates.

Let’s look deeper into how we can implement some scenarios leveraging multiple conditions to generate count queries. Suppose we want to get a count of all the books containing a certain keyword in the title:

long countBooksByTitle(String titleKeyword) {
    CriteriaBuilder cb = entityManager.getCriteriaBuilder();
    CriteriaQuery<Long> cq = cb.createQuery(Long.class);
    Root<Book> bookRoot = cq.from(Book.class);
    Predicate condition = cb.like(bookRoot.get("title"), "%" +    titleKeyword + "%");
    cq.where(condition);
    cq.select(cb.count(bookRoot));
    return entityManager.createQuery(cq).getSingleResult()
}

In addition to previous steps, for the conditional count, we create a Predicate that represents the WHERE clause of the SQL query.

The cb.like method creates a condition that checks if the title of the book contains the title keyword. The  % is the wildcard that matches any sequence of characters. We then add this predicate to the CriteriaQuery using cq.where(condition), which applies this condition to the query.

Another use case in our scenario may be fetching a count of all the books by an author:

long countBooksByAuthor(String authorName) {        
    CriteriaBuilder cb = entityManager.getCriteriaBuilder();
    CriteriaQuery<Long> cq = cb.createQuery(Long.class);
    Root<Book> bookRoot = cq.from(Book.class);
    Predicate condition = cb.equal(bookRoot.get("author"), authorName);
    cq.where(condition);
    cq.select(cb.count(bookRoot));
    return entityManager.createQuery(cq).getSingleResult();
}

Here, the predicate is based on the cb.equal() method which only filters the records containing the exact author name.

5. Combining Multiple Criteria

Criteria Query allows us to combine multiple criteria using logical operators such as AND, OR, and NOT. Let’s consider some cases where we want a count of books based on multiple criteria.  Suppose, we want to get all the book counts with certain authors, titles, and publishing years:

long countBooksByAuthorOrYear(int publishYear, String authorName) {
    CriteriaBuilder cb = entityManager.getCriteriaBuilder();
    CriteriaQuery<Long> cq = cb.createQuery(Long.class);
    Root<Book> bookRoot = cq.from(Book.class);
    Predicate authorCondition = cb.equal(bookRoot.get("author"), authorName);
    Predicate yearCondition = cb.greaterThanOrEqualTo(bookRoot.get("publishYear"), 1800);
    cq.where(cb.or(authorCondition, yearCondition));
    cq.select(cb.count(bookRoot));
    return entityManager.createQuery(cq).getSingleResult();
}

Here, we create two predicates representing conditions on the author and publishing year of the books. We then combine these predicates using cb.or() to form a compound condition.

Similarly, we can also have a scenario where we want to get the counts of books that have a certain title or they have a combination of author and year:

long countBooksByTitleOrYearAndAuthor(String authorName, int publishYear, String titleKeyword) {
    CriteriaBuilder cb = entityManager.getCriteriaBuilder();
    CriteriaQuery<Long> cq = cb.createQuery(Long.class);
    Root<Book> bookRoot = cq.from(Book.class);

    Predicate authorCondition = cb.equal(bookRoot.get("author"), authorName);
    Predicate yearCondition = cb.equal(bookRoot.get("publishYear"), publishYear);
    Predicate titleCondition = cb.like(bookRoot.get("title"), "%" + titleKeyword + "%");

    Predicate authorAndYear = cb.and(authorCondition, yearCondition);
    cq.where(cb.or(authorAndYear, titleCondition));
    cq.select(cb.count(bookRoot));

    return entityManager.createQuery(cq).getSingleResult();
}

We again created three predicates, but now we want to have an or between authorAndYearCondition predicate and titleCondition predicate using cb.or(authorAndYear, titleCondition)

6. Integration Tests

Let’s now provide a pattern for performing integration tests ensuring our count queries work. We’ll use the @DataJPATest  annotation provided by Spring to inject the necessary Repository layer in our tests using H2 in the memory Database as an underlying persistence storage. We’ll inject TestEntityManager in our test class and use it to insert data. Let’s take one case of getting the count of all the books by a certain author:

@Test
void givenBookDataAdded_whenCountBooksByAuthor_thenReturnsCount() {
    entityManager.persist(new Book("Java Book 1", "Author 1", 1967, "Non Fiction"));
    entityManager.persist(new Book("Java Book 2", "Author 1", 1999, "Non Fiction"));
    entityManager.persist(new Book("Spring Book", "Author 2", 2007, "Non Fiction"));

    long count = bookRepository.countBooksByAuthor("Author 1");

    assertEquals(2, count);
}

Similar to the above, we can write tests for all the different count scenarios provided in the repository.

7. Conclusion

In this article, we demonstrated how to perform conditional counting in a Spring Boot application using the JPA Criteria API. We set up a basic library management system and implemented a custom repository method to count books based on specific conditions.