hibernate-demo/src/test/java/com/ankurm/hibernatedemo/association/NPlusOneTest.java

package com.ankurm.hibernatedemo.association;

import static org.assertj.core.api.Assertions.assertThat;

import jakarta.persistence.EntityGraph;
import jakarta.persistence.EntityManager;
import jakarta.persistence.EntityManagerFactory;
import jakarta.persistence.TypedQuery;
import java.util.List;
import java.util.Map;
import org.hibernate.SessionFactory;
import org.hibernate.stat.Statistics;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.transaction.annotation.Transactional;

/**
 * Backs ankurm.com post 4872, docs/12-association-mappings.md chapter "Counting the N+1".
 *
 * <p>Seeds 100 authors with 3 books each (301 rows total incl. authors) and counts real
 * {@code Statistics.getPrepareStatementCount()} for four strategies: naive lazy iteration,
 * JPQL fetch join, {@code @EntityGraph}, and {@code @BatchSize(10)}. Run with
 * {@code mvn -Dtest=NPlusOneTest test}.
 */
@SpringBootTest
class NPlusOneTest {

    private static final Logger DEMO = LoggerFactory.getLogger("DEMO");
    private static final int AUTHOR_COUNT = 100;
    private static final int BOOKS_PER_AUTHOR = 3;

    @Autowired
    private EntityManagerFactory emf;

    @BeforeEach
    void cleanTables() {
        EntityManager em = emf.createEntityManager();
        em.getTransaction().begin();
        em.createQuery("DELETE FROM AssocBook").executeUpdate();
        em.createQuery("DELETE FROM AssocAuthor").executeUpdate();
        em.createQuery("DELETE FROM BatchBook").executeUpdate();
        em.createQuery("DELETE FROM BatchAuthor").executeUpdate();
        em.getTransaction().commit();
        em.close();
    }

    private Statistics stats() {
        return emf.unwrap(SessionFactory.class).getStatistics();
    }

    private void seedPlainAuthors() {
        EntityManager em = emf.createEntityManager();
        em.getTransaction().begin();
        for (int a = 0; a < AUTHOR_COUNT; a++) {
            AssocAuthor author = new AssocAuthor("Author-" + a);
            em.persist(author);
            for (int b = 0; b < BOOKS_PER_AUTHOR; b++) {
                em.persist(new AssocBook("Book-" + a + "-" + b, author));
            }
        }
        em.getTransaction().commit();
        em.close();
    }

    private void seedBatchAuthors() {
        EntityManager em = emf.createEntityManager();
        em.getTransaction().begin();
        for (int a = 0; a < AUTHOR_COUNT; a++) {
            BatchAuthor author = new BatchAuthor("BatchAuthor-" + a);
            em.persist(author);
            for (int b = 0; b < BOOKS_PER_AUTHOR; b++) {
                em.persist(new BatchBook("BatchBook-" + a + "-" + b, author));
            }
        }
        em.getTransaction().commit();
        em.close();
    }

    @Test
    void naiveLazyIteration_firesOneQueryPerAuthor_theClassicNPlusOne() {
        seedPlainAuthors();
        EntityManager em = emf.createEntityManager();
        stats().clear();

        List<AssocAuthor> authors = em.createQuery("SELECT a FROM AssocAuthor a", AssocAuthor.class).getResultList();
        long booksTouched = 0;
        for (AssocAuthor author : authors) {
            booksTouched += author.getBooks().size(); // triggers one SELECT per author
        }

        long queries = stats().getPrepareStatementCount();
        DEMO.info("NAIVE lazy iteration: {} authors, {} queries (1 for authors + {} for their book collections), books touched = {}",
                authors.size(), queries, authors.size(), booksTouched);

        assertThat(authors).hasSize(AUTHOR_COUNT);
        assertThat(queries)
                .as("1 query to load authors + 1 per author to lazily load its books == N+1")
                .isEqualTo(1 + AUTHOR_COUNT);
        em.close();
    }

    @Test
    void jpqlFetchJoin_firesExactlyOneQuery() {
        seedPlainAuthors();
        EntityManager em = emf.createEntityManager();
        stats().clear();

        List<AssocAuthor> authors = em.createQuery(
                        "SELECT DISTINCT a FROM AssocAuthor a JOIN FETCH a.books", AssocAuthor.class)
                .getResultList();
        long booksTouched = 0;
        for (AssocAuthor author : authors) {
            booksTouched += author.getBooks().size(); // already initialised, no extra SELECT
        }

        long queries = stats().getPrepareStatementCount();
        DEMO.info("JPQL JOIN FETCH: {} authors, {} queries, books touched = {}", authors.size(), queries, booksTouched);

        assertThat(authors).hasSize(AUTHOR_COUNT);
        assertThat(queries).as("fetch join collapses N+1 into a single query").isEqualTo(1);
        em.close();
    }

    @Test
    void entityGraph_firesExactlyOneQuery() {
        seedPlainAuthors();
        EntityManager em = emf.createEntityManager();
        stats().clear();

        EntityGraph<?> graph = em.getEntityGraph("AssocAuthor.books");
        TypedQuery<AssocAuthor> query = em.createQuery("SELECT a FROM AssocAuthor a", AssocAuthor.class);
        query.setHint("jakarta.persistence.fetchgraph", graph);
        List<AssocAuthor> authors = query.getResultList();
        long booksTouched = 0;
        for (AssocAuthor author : authors) {
            booksTouched += author.getBooks().size();
        }

        long queries = stats().getPrepareStatementCount();
        DEMO.info("@EntityGraph (fetchgraph hint): {} authors, {} queries, books touched = {}", authors.size(), queries, booksTouched);

        assertThat(authors).hasSize(AUTHOR_COUNT);
        assertThat(queries).as("entity graph fetch also collapses to a single query").isEqualTo(1);
        em.close();
    }

    @Test
    void batchSize10_collapsesNPlusOneIntoCeilNOverBatchSizePlusOne() {
        seedBatchAuthors();
        EntityManager em = emf.createEntityManager();
        stats().clear();

        List<BatchAuthor> authors = em.createQuery("SELECT a FROM BatchAuthor a", BatchAuthor.class).getResultList();
        long booksTouched = 0;
        for (BatchAuthor author : authors) {
            booksTouched += author.getBooks().size(); // triggers 1 SELECT per 10 authors (batch fetch)
        }

        long queries = stats().getPrepareStatementCount();
        // ceil(100 / 10) = 10 batch selects + 1 for the authors themselves = 11
        long expected = 1 + (long) Math.ceil(AUTHOR_COUNT / 10.0);
        DEMO.info("@BatchSize(10): {} authors, {} queries (expected {} = 1 + ceil({}/10)), books touched = {}",
                authors.size(), queries, expected, AUTHOR_COUNT, booksTouched);

        assertThat(authors).hasSize(AUTHOR_COUNT);
        assertThat(queries).as("BatchSize(10) turns 101 queries into ceil(100/10)+1 = 11").isEqualTo(expected);
        em.close();
    }

    @Test
    void printSideBySideComparisonTable() {
        // Re-derive each number in one place so the article table and the assertions above can
        // never silently drift apart.
        seedPlainAuthors();
        seedBatchAuthors();

        EntityManager em = emf.createEntityManager();

        stats().clear();
        List<AssocAuthor> naive = em.createQuery("SELECT a FROM AssocAuthor a", AssocAuthor.class).getResultList();
        naive.forEach(a -> a.getBooks().size());
        long naiveQ = stats().getPrepareStatementCount();

        stats().clear();
        List<AssocAuthor> fetchJoin = em.createQuery(
                        "SELECT DISTINCT a FROM AssocAuthor a JOIN FETCH a.books", AssocAuthor.class)
                .getResultList();
        fetchJoin.forEach(a -> a.getBooks().size());
        long fetchJoinQ = stats().getPrepareStatementCount();

        stats().clear();
        EntityGraph<?> graph = em.getEntityGraph("AssocAuthor.books");
        TypedQuery<AssocAuthor> egQuery = em.createQuery("SELECT a FROM AssocAuthor a", AssocAuthor.class);
        egQuery.setHint("jakarta.persistence.fetchgraph", graph);
        List<AssocAuthor> eg = egQuery.getResultList();
        eg.forEach(a -> a.getBooks().size());
        long egQ = stats().getPrepareStatementCount();

        stats().clear();
        List<BatchAuthor> batched = em.createQuery("SELECT a FROM BatchAuthor a", BatchAuthor.class).getResultList();
        batched.forEach(a -> a.getBooks().size());
        long batchQ = stats().getPrepareStatementCount();

        DEMO.info("=== Side-by-side query counts for {} authors x {} books each ===", AUTHOR_COUNT, BOOKS_PER_AUTHOR);
        DEMO.info("naive lazy iteration : {} queries", naiveQ);
        DEMO.info("JPQL JOIN FETCH      : {} queries", fetchJoinQ);
        DEMO.info("@EntityGraph         : {} queries", egQ);
        DEMO.info("@BatchSize(10)       : {} queries", batchQ);

        assertThat(Map.of("naive", naiveQ, "fetchJoin", fetchJoinQ, "entityGraph", egQ, "batch10", batchQ))
                .isEqualTo(Map.of("naive", 101L, "fetchJoin", 1L, "entityGraph", 1L, "batch10", 11L));
        em.close();
    }
}