How to detect the Hibernate N+1 query problem during testing

Imagine having a tool that can automatically detect if you are using JPA and Hibernate properly. Hypersistence Optimizer is that tool!

Introduction

In this article, you are going to learn how to automatically detect the N+1 query problem when using JPA and Hibernate using the db-util open-source project.

With Hibernate, you manage entity state transitions which are then translated to SQL statements. The number of generated SQL statements is affected by the current fetching strategy, Criteria queries or Collection mappings and you might not always get what you expected. Ignoring SQL statements is risky and it may eventually put a heavy toll on the overall application performance.

I’m a strong advocate of peer review, but that’s not “sine qua non” for detecting bad Hibernate usage. Subtle changes may affect the SQL statement count and pass unnoticed through the reviewing process. Not in the least, when it comes to “guessing” the JPA SQL statements, I feel like I can use any extra help. I’m for as much automation as possible, and that’s why I came up with a mechanism for enforcing the SQL statement count expectations.

First, we need a way to intercept all executed SQL statements. I researched on this topic and I was lucky to find this great datasource-proxy library.

Adding an automated validator

This safeguard is meant to run only during the testing phase, so I’ll add it to the Integration Testing Spring context exclusively. I’ve already talked about Spring bean aliasing and this is the right time to make use of it.

@Bean
public DataSource dataSource(DataSource originalDataSource) {
    ChainListener listener = new ChainListener();
    SLF4JQueryLoggingListener loggingListener = new SLF4JQueryLoggingListener();
    loggingListener.setQueryLogEntryCreator(new InlineQueryLogEntryCreator());
    listener.addListener(loggingListener);
    listener.addListener(new DataSourceQueryCountListener());
    return ProxyDataSourceBuilder
        .create(originalDataSource)
        .name("DS-Proxy")
        .listener(listener)
        .build();
}

The new proxy data source decorates the already existing data source, intercepting all executed SQL statements. This library can log all SQL statements along with the actual parameter values, unlike the default Hibernate logging which only prints a placeholder instead.

This is how the validator looks like:

public class SQLStatementCountValidator {

    private SQLStatementCountValidator() {
    }

    /**
     * Reset the statement recorder
     */
    public static void reset() {
        QueryCountHolder.clear();
    }

    /**
     * Assert select statement count
     * @param expectedSelectCount expected select statement count
     */
    public static void assertSelectCount(int expectedSelectCount) {
        QueryCount queryCount = QueryCountHolder.getGrandTotal();
        int recordedSelectCount = queryCount.getSelect();
        if(expectedSelectCount != recordedSelectCount) {
            throw new SQLSelectCountMismatchException(
				expectedSelectCount, 
				recordedSelectCount
			);
        }
    }

    /**
     * Assert insert statement count
     * @param expectedInsertCount expected insert statement count
     */
    public static void assertInsertCount(int expectedInsertCount) {
        QueryCount queryCount = QueryCountHolder.getGrandTotal();
        int recordedInsertCount = queryCount.getInsert();
        if(expectedInsertCount != recordedInsertCount) {
            throw new SQLInsertCountMismatchException(
				expectedInsertCount, 
				recordedInsertCount
			);
        }
    }

    /**
     * Assert update statement count
     * @param expectedUpdateCount expected update statement count
     */
    public static void assertUpdateCount(int expectedUpdateCount) {
        QueryCount queryCount = QueryCountHolder.getGrandTotal();
        int recordedUpdateCount = queryCount.getUpdate();
        if(expectedUpdateCount != recordedUpdateCount) {
            throw new SQLUpdateCountMismatchException(
				expectedUpdateCount, 
				recordedUpdateCount
			);
        }
    }

    /**
     * Assert delete statement count
     * @param expectedDeleteCount expected delete statement count
     */
    public static void assertDeleteCount(int expectedDeleteCount) {
        QueryCount queryCount = QueryCountHolder.getGrandTotal();
        int recordedDeleteCount = queryCount.getDelete();
        if(expectedDeleteCount != recordedDeleteCount) {
            throw new SQLDeleteCountMismatchException(
				expectedDeleteCount, 
				recordedDeleteCount
			);
        }
    }
}

This utility is part of my db-util project along with the JPA and MongoDB optimistic concurrency control retry mechanism.

Since it’s already available in Maven Central Repository, you can easily use it by just adding this dependency to your pom.xml:

<dependency>
    <groupId>com.vladmihalcea</groupId>
    <artifactId>db-util</artifactId>
    <version>${db-util.version}</version>
</dependency>

Let’s write a test for detecting the infamous N+1 select query problem.

For this we will write two service methods, one of them being affected by the aforementioned issue:

@Override
@Transactional
public List<WarehouseProductInfo> findAllWithNPlusOne() {
    List<WarehouseProductInfo> warehouseProductInfos = entityManager
    .createQuery(
        "from WarehouseProductInfo", WarehouseProductInfo.class)
    .getResultList();
	
    navigateWarehouseProductInfos(warehouseProductInfos);
	
    return warehouseProductInfos;
}

@Override
@Transactional
public List<WarehouseProductInfo> findAllWithFetch() {
    List<WarehouseProductInfo> warehouseProductInfos = entityManager
    .createQuery(
        "from WarehouseProductInfo wpi " +
        "join fetch wpi.product p " +
        "join fetch p.company", WarehouseProductInfo.class)
    .getResultList();
	
    navigateWarehouseProductInfos(warehouseProductInfos);
    
	return warehouseProductInfos;
}

private void navigateWarehouseProductInfos(
    List<WarehouseProductInfo> warehouseProductInfos) {
    for(WarehouseProductInfo warehouseProductInfo : warehouseProductInfos) {
        warehouseProductInfo.getProduct();
    }
}

The unit test is rather simple to write since it follows the same coding style of any other JUnit assert mechanism.

try {
    SQLStatementCountValidator.reset();
    warehouseProductInfoService.findAllWithNPlusOne();
    assertSelectCount(1);
} catch (SQLSelectCountMismatchException e) {
    assertEquals(3, e.getRecorded());
}

SQLStatementCountValidator.reset();
warehouseProductInfoService.findAllWithFetch();
assertSelectCount(1);

Our validator works for all SQL statement types, so let’s check how many SQL INSERTs are being executed by the following service method:

@Override
@Transactional
public WarehouseProductInfo newWarehouseProductInfo() {

    LOGGER.info("newWarehouseProductInfo");

    Company company = entityManager
    .createQuery("from Company", Company.class)
    .getResultList()
    .get(0);

    Product product3 = new Product("phoneCode");
    product3.setName("Phone");
    product3.setCompany(company);

    WarehouseProductInfo warehouseProductInfo3 = new WarehouseProductInfo();
    warehouseProductInfo3.setQuantity(19);
    product3.addWarehouse(warehouseProductInfo3);

    entityManager.persist(product3);
    return warehouseProductInfo3;
}

And the validator looks like:

SQLStatementCountValidator.reset();

warehouseProductInfoService.newWarehouseProductInfo();

assertSelectCount(1);
assertInsertCount(2);

Let’s check the test logs to convince ourselves of its effectiveness:

-- newWarehouseProductInfo

SELECT 
	c.id as id1_6_, 
	c.name as name2_6_ 
FROM 
	Company c

INSERT INTO WarehouseProductInfo 
	(id, quantity) 
VALUES 
	(default, 19)
	
INSERT INTO Product 
	(id, code, company_id, importer_id, name, version) 
VALUES 
	(default, 'phoneCode', 1, -5, 'Phone', 0)

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

Code reviewing is a fine technique, but it’s not enough for a large-scale development project. That’s why automatic checking is of paramount importance. Once the test was written, you are assured that no future change can break your assumptions.

Code available on GitHub.

Transactions and Concurrency Control eBook

8 Comments on “How to detect the Hibernate N+1 query problem during testing

  1. Db-util is a great tool. If you use it with Spring Boot, like I do, maybe the next observations are of any value to you:

    1: the integration test

    Maybe your Spring boot integration test is annotated something like:

    @SpringBootTest(classes = Application.class)
    @RunWith(SpringRunner.class)
    @ActiveProfiles("test")
    @Transactional
    

    The transactional part is nice: all modifications are rolled back. However, the transaction will be propagated to the method under test. Updates in method under test are postponed by Hibernate because the wrapping transaction is not closed yet. To simulate the end of the transaction in the method under test, you should flush te persistence context:

    @PersistenceContext
    private EntityManager entityManager;
    
    @Autowired
    private AnswerService myService;
    
    @Test
    public void testBatchedUpdates(){
        SQLStatementCountValidator.reset();
        myService.update5Entities();
        entityManager.flush(); // Without this, no updates
        assertUpdateCount(1);
    }
    

    Remark that it is flushed, but not committed, so after this test has run all data is cleared

    2: circular dependency issue with datasource proxy

    Comment of Haider on 24 april suggests a solution. Based on this solution, a simpler one worked for me too:

    @Component
    public class DatasourceProxyBeanPostProcessor implements BeanPostProcessor {
    
    @Override
    public Object postProcessAfterInitialization(final Object bean, final String beanName) throws BeansException {
        if (bean instanceof DataSource) {
            ChainListener listener = new ChainListener();
            listener.addListener(new DataSourceQueryCountListener());
            return ProxyDataSourceBuilder
                    .create((DataSource) bean)
                    .listener(listener)
                    .build();
        }
        return bean;
    }
    
    }
    

    3: using it in combination with the flexipool starter

    Why only use 1 great tool of Vlad? Flexipool is anther great tool and it has a starter for spring boot. Using both this starter and db-util requires version 1.0.5 for db-util.

  2. SQLStatementCountValidator.reset();
    warehouseProductInfoService.newWarehouseProductInfo();
    assertSelectCount(1);
    Hello. When I run the code. An error occurs all the time. recorderSelectCount equal to zero. What am I doing wrong?

    • The code is on Github and works like a charm. You need to debug my test case and yours and see there the execution differs.

  3. Did I misread something, or there is no definition in your article what “N+1 problem” actually is?

    First time you use this term is: “Let’s write a test for detecting the infamous N+1 select query problem.” …

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.