spring-boot-demo/graphql-dataloader/src/test/java/com/ankurm/graphqldataloader/BatchedResolverSqlLogTest.java

package com.ankurm.graphqldataloader;

import com.ankurm.graphqldataloader.domain.Author;
import com.ankurm.graphqldataloader.domain.AuthorRepository;
import com.ankurm.graphqldataloader.domain.Book;
import com.ankurm.graphqldataloader.domain.BookRepository;
import com.ankurm.graphqldataloader.support.SqlLog;
import tools.jackson.databind.JsonNode;
import tools.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.MethodOrderer;
import org.junit.jupiter.api.Order;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestMethodOrder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.web.server.LocalServerPort;
import org.springframework.test.context.ActiveProfiles;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardOpenOption;
import java.util.ArrayList;
import java.util.List;

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

/**
 * Same experiments as {@link NaiveResolverSqlLogTest}, run against the {@code batched} profile
 * ({@link com.ankurm.graphqldataloader.resolver.BatchedAuthorResolver}) instead — same schema,
 * same seed helpers, same GraphQL query text, only the resolver wiring differs. The SQL statement
 * counts are what the naive-vs-batched comparison in the article is built on.
 */
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
@ActiveProfiles("batched")
@TestMethodOrder(MethodOrderer.OrderAnnotation.class)
class BatchedResolverSqlLogTest {

    @LocalServerPort
    int port;

    @Autowired
    AuthorRepository authorRepository;

    @Autowired
    BookRepository bookRepository;

    @Autowired
    SqlLog sqlLog;

    private final HttpClient http = HttpClient.newHttpClient();
    private final ObjectMapper mapper = new ObjectMapper();

    @BeforeEach
    void cleanDatabase() {
        bookRepository.deleteAll();
        authorRepository.deleteAll();
    }

    @Test
    @Order(1)
    void a_fiveBooksFromFiveDistinctAuthorsIssueOneBatchedAuthorQuery() throws Exception {
        seedDistinctAuthorPerBook(5);
        sqlLog.reset();

        JsonNode data = postGraphql("{ books { title author { name } } }");

        assertThat(data.get("books")).hasSize(5);
        for (JsonNode book : data.get("books")) {
            assertThat(book.get("author").get("name").asText()).isNotBlank();
        }
        // 1 SELECT for books.findAll() + 1 SELECT ... WHERE id IN (...) from the batched
        // authorRepository.findAllById() call — GraphQL Java's DataLoader collected all 5 pending
        // author lookups into a single batch before BatchedAuthorResolver.author() ever ran.
        assertThat(sqlLog.count()).isEqualTo(2);
        assertThat(sqlLog.all().get(1)).containsIgnoringCase(" in (");

        writeTranscript("batched-a-five-distinct-authors",
                "batched profile / 5 books, 5 distinct authors (no repeats to dedupe)");
    }

    @Test
    @Order(2)
    void b_twentyBooksFromFiveAuthorsDedupesToOneBatchedQueryForFiveAuthors() throws Exception {
        seedRoundRobinAuthors(5, 20);
        sqlLog.reset();

        JsonNode data = postGraphql("{ books { title author { name } } }");

        assertThat(data.get("books")).hasSize(20);
        // Still exactly 2 statements: 1 for books.findAll(), 1 for the batched author lookup — the
        // resolver's `.distinct()` call collapses the 20 pending DataLoader keys down to the 5
        // unique authorIds actually present before the repository is ever asked.
        assertThat(sqlLog.count()).isEqualTo(2);

        long placeholders = sqlLog.all().get(1).chars().filter(c -> c == '?').count();
        assertThat(placeholders).isEqualTo(5);

        writeTranscript("batched-b-twenty-books-five-authors",
                "batched profile / 20 books, 5 distinct authors (4 books per author) — one batched IN query");
    }

    @Test
    @Order(3)
    void c_danglingAuthorIdNullsTheEntireBooksListViaNonNullPropagationSameAsNaive() throws Exception {
        Author real = authorRepository.save(new Author("Ursula K. Le Guin"));
        bookRepository.save(new Book("The Dispossessed", real.getId()));
        bookRepository.save(new Book("Orphan Book With No Author Row", 9_999_999L));
        sqlLog.reset();

        String responseBody = postGraphqlRaw("{ books { title author { name } } }");
        JsonNode root = mapper.readTree(responseBody);

        // Same outcome as the naive profile: BatchedAuthorResolver.author() simply omits the map
        // entry for the orphan Book, and Spring GraphQL treats a missing DataLoader/batch result
        // the same as an explicit null for that source object — the non-null `author: Author!`
        // field then propagates null up through the non-null `[Book!]!` list to the whole response.
        assertThat(root.get("data").isNull()).isTrue();
        assertThat(root.get("errors")).isNotEmpty();
        assertThat(root.get("errors").get(0).get("message").asText())
                .contains("non-null");

        Files.writeString(
                outputPath("batched-c-dangling-foreign-key-null-propagation.txt"),
                "batched profile / one Book has authorId=9999999 which matches no Author row\n"
                        + "raw HTTP response body from POST /graphql:\n\n" + responseBody + "\n",
                StandardCharsets.UTF_8, StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
    }

    private void seedDistinctAuthorPerBook(int n) {
        for (int i = 1; i <= n; i++) {
            Author author = authorRepository.save(new Author("Author " + i));
            bookRepository.save(new Book("Book " + i, author.getId()));
        }
    }

    private void seedRoundRobinAuthors(int authorCount, int bookCount) {
        List<Long> authorIds = new ArrayList<>();
        for (int i = 1; i <= authorCount; i++) {
            authorIds.add(authorRepository.save(new Author("Author " + i)).getId());
        }
        for (int i = 1; i <= bookCount; i++) {
            Long authorId = authorIds.get((i - 1) % authorCount);
            bookRepository.save(new Book("Book " + i, authorId));
        }
    }

    private JsonNode postGraphql(String query) throws Exception {
        String body = postGraphqlRaw(query);
        JsonNode root = mapper.readTree(body);
        assertThat(root.has("errors")).isFalse();
        return root.get("data");
    }

    private String postGraphqlRaw(String query) throws Exception {
        String requestBody = mapper.writeValueAsString(java.util.Map.of("query", query));
        HttpRequest request = HttpRequest.newBuilder()
                .uri(URI.create("http://localhost:" + port + "/graphql"))
                .header("Content-Type", "application/json")
                .header("Accept", "application/json")
                .POST(HttpRequest.BodyPublishers.ofString(requestBody, StandardCharsets.UTF_8))
                .build();
        HttpResponse<String> response = http.send(request, HttpResponse.BodyHandlers.ofString());
        assertThat(response.statusCode()).isEqualTo(200);
        return response.body();
    }

    private void writeTranscript(String fileNameStem, String heading) throws Exception {
        String content = heading + "\n\n" + sqlLog.render() + "\n";
        Files.writeString(outputPath(fileNameStem + ".txt"), content, StandardCharsets.UTF_8,
                StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING);
    }

    private Path outputPath(String fileName) throws Exception {
        Path dir = Path.of(System.getProperty("user.dir"), "docs", "output");
        Files.createDirectories(dir);
        return dir.resolve(fileName);
    }
}