JPA Lazy Loading and the N+1 Problem in Spring Boot

spring.jpa.properties.hibernate.generate_statistics=true

//Publisher.java
package com.javabrahman.data.lazy;

import jakarta.persistence.CascadeType;
import jakarta.persistence.Entity;
import jakarta.persistence.GeneratedValue;
import jakarta.persistence.GenerationType;
import jakarta.persistence.Id;
import jakarta.persistence.OneToMany;
import java.util.ArrayList;
import java.util.List;

@Entity
public class Publisher {
  @Id
  @GeneratedValue(strategy = GenerationType.IDENTITY)
  private Long id;
  private String name;

  @OneToMany(mappedBy = "publisher", cascade = CascadeType.ALL)   // lazy by default
  private List<Book> books = new ArrayList<>();

  protected Publisher() {
  }

  public Publisher(String name) {
    this.name = name;
  }

  public void addBook(Book book) {
    books.add(book);
    book.setPublisher(this);
  }

  public String getName() { return name; }
  public List<Book> getBooks() { return books; }
}

//Book.java
package com.javabrahman.data.lazy;

import jakarta.persistence.Entity;
import jakarta.persistence.FetchType;
import jakarta.persistence.GeneratedValue;
import jakarta.persistence.GenerationType;
import jakarta.persistence.Id;
import jakarta.persistence.ManyToOne;

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

  @ManyToOne(fetch = FetchType.LAZY)
  private Publisher publisher;

  protected Book() {
  }

  public Book(String title) {
    this.title = title;
  }

  public String getTitle() { return title; }
  public Publisher getPublisher() { return publisher; }
  void setPublisher(Publisher publisher) { this.publisher = publisher; }
}

//PublisherRepository.java
package com.javabrahman.data.lazy;

import java.util.List;
import org.springframework.data.domain.Page;
import org.springframework.data.domain.Pageable;
import org.springframework.data.jpa.repository.EntityGraph;
import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.Query;

public interface PublisherRepository extends JpaRepository<Publisher, Long> {

  // Fix 1: join fetch loads each publisher together with its books, in one query
  @Query("select distinct p from Publisher p join fetch p.books")
  List<Publisher> findAllWithBooks();

  // Fix 2: an entity graph says the same thing without touching the query
  @EntityGraph(attributePaths = "books")
  List<Publisher> findAllBy();

  // Fetching a collection and paging do not mix (see below)
  @Query(value = "select distinct p from Publisher p join fetch p.books", countQuery = "select count(p) from Publisher p")
  Page<Publisher> findPageWithBooks(Pageable pageable);
}

//LazyApp.java
package com.javabrahman.data.lazy;

import jakarta.persistence.EntityManagerFactory;
import java.util.List;
import java.util.function.Supplier;
import org.hibernate.LazyInitializationException;
import org.hibernate.SessionFactory;
import org.hibernate.stat.Statistics;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.data.domain.PageRequest;
import org.springframework.transaction.support.TransactionTemplate;

@SpringBootApplication
public class LazyApp {

  @Bean
  CommandLineRunner demo(PublisherRepository publishers, EntityManagerFactory emf, TransactionTemplate tx) {
    Statistics stats = emf.unwrap(SessionFactory.class).getStatistics();

    return args -> {
      for (String name : List.of("Addison-Wesley", "Manning", "O'Reilly", "Pragmatic Bookshelf")) {
        Publisher p = new Publisher(name);
        for (int i = 1; i <= 3; i++) {
          p.addBook(new Book(name + " book " + i));
        }
        publishers.save(p);
      }

      System.out.println("== lazy loading outside a transaction");
      Publisher detached = publishers.findAll().get(0);
      try {
        System.out.println(detached.getBooks().size());
      } catch (LazyInitializationException e) {
        System.out.println("LazyInitializationException: " + e.getMessage());
      }

      System.out.println("== lazy loading inside a transaction");
      count(stats, "load one publisher and read its books", () -> tx.execute(status -> {
        Publisher p = publishers.findAll().get(0);
        return p.getBooks().size();
      }));

      System.out.println("== N+1 loop");
      count(stats, "N+1 loop: publishers and their books", () -> tx.execute(status -> {
        int books = 0;
        for (Publisher p : publishers.findAll()) {   // 1 query for the publishers ...
          books += p.getBooks().size();              // ... and 1 more for each publisher's books
        }
        return books;
      }));

      System.out.println("== fix 1: join fetch");
      count(stats, "join fetch", () -> tx.execute(status -> publishers.findAllWithBooks().stream().mapToInt(p -> p.getBooks().size()).sum()));

      System.out.println("== fix 2: entity graph");
      count(stats, "entity graph", () -> tx.execute(status -> publishers.findAllBy().stream().mapToInt(p -> p.getBooks().size()).sum()));

      System.out.println("== paging a fetched collection");
      count(stats, "page of publishers with join fetch", () -> tx.execute(status -> publishers.findPageWithBooks(PageRequest.of(0, 2)).getContent().size()));
    };
  }

  // Runs the work and prints how many SQL statements it prepared
  private static void count(Statistics stats, String label, Supplier<Object> work) {
    stats.clear();
    Object result = work.get();
    System.out.println(label + ": result " + result + ", SQL statements: " + stats.getPrepareStatementCount());
  }

  public static void main(String[] args) {
    SpringApplication.run(LazyApp.class, args);
  }
}