BAEL-9323

spring-batch-2/pom.xml

@@ -104,6 +104,11 @@
                 </exclusion>
             </exclusions>
         </dependency>
+        <dependency>
+            <groupId>com.graphql-java</groupId>
+            <artifactId>java-dataloader</artifactId>
+            <version>3.2.0</version>
+        </dependency>
     </dependencies>
 
     <build>

spring-batch-2/src/main/java/com/baeldung/dataloaderbatchprocessing/DataLoaderApp.java

@@ -0,0 +1,13 @@
+package com.baeldung.dataloaderbatchprocessing;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication
+public class DataLoaderApp {
+
+    public static void main(String[] args) {
+        SpringApplication.run(DataLoaderApp.class, args);
+    }
+}
+

spring-batch-2/src/main/java/com/baeldung/dataloaderbatchprocessing/UserDataLoader.java

@@ -0,0 +1,37 @@
+package com.baeldung.dataloaderbatchprocessing;
+
+import org.dataloader.BatchLoader;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderFactory;
+import org.springframework.stereotype.Component;
+
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.baeldung.dataloaderbatchprocessing.entity.User;
+import com.baeldung.dataloaderbatchprocessing.service.UserService;
+
+@Component
+public class UserDataLoader {
+
+    private final UserService userService;
+
+    public UserDataLoader(UserService userService) {
+        this.userService = userService;
+    }
+
+    public DataLoader<String, User> createUserLoader() {
+        BatchLoader<String, User> userBatchLoader = ids -> {
+            return userService.getUsersByIds(ids)
+                .thenApply(users -> {
+                    Map<String, User> userMap = users.stream()
+                        .collect(Collectors.toMap(User::getId, user -> user));
+                    return ids.stream()
+                        .map(userMap::get)
+                        .collect(Collectors.toList());
+                });
+        };
+
+        return DataLoaderFactory.newDataLoader(userBatchLoader);
+    }
+}

spring-batch-2/src/main/java/com/baeldung/dataloaderbatchprocessing/entity/User.java

@@ -0,0 +1,19 @@
+package com.baeldung.dataloaderbatchprocessing.entity;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+
+@Entity
+@Getter
+@AllArgsConstructor
+@NoArgsConstructor
+public class User {
+
+    @Id
+    private String id;
+    private String name;
+}
+

spring-batch-2/src/main/java/com/baeldung/dataloaderbatchprocessing/repository/UserRepository.java

@@ -0,0 +1,9 @@
+package com.baeldung.dataloaderbatchprocessing.repository;
+
+import org.springframework.data.jpa.repository.JpaRepository;
+
+import com.baeldung.dataloaderbatchprocessing.entity.User;
+
+public interface UserRepository extends JpaRepository<User, String> {
+}
+

spring-batch-2/src/main/java/com/baeldung/dataloaderbatchprocessing/service/UserService.java

@@ -0,0 +1,22 @@
+package com.baeldung.dataloaderbatchprocessing.service;
+import org.springframework.stereotype.Service;
+
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import com.baeldung.dataloaderbatchprocessing.entity.User;
+import com.baeldung.dataloaderbatchprocessing.repository.UserRepository;
+
+@Service
+public class UserService {
+    private final UserRepository userRepository;
+
+    public UserService(UserRepository userRepository) {
+        this.userRepository = userRepository;
+    }
+
+    public CompletableFuture<List<User>> getUsersByIds(List<String> ids) {
+        return CompletableFuture.supplyAsync(() -> userRepository.findAllById(ids));
+    }
+}
+

spring-batch-2/src/test/java/com/baeldung/dataloaderbatchprocessing/DataLoaderUnitTest.java

@@ -0,0 +1,57 @@
+package com.baeldung.dataloaderbatchprocessing;
+
+import com.baeldung.dataloaderbatchprocessing.entity.User;
+import com.baeldung.dataloaderbatchprocessing.service.UserService;
+
+import org.dataloader.DataLoader;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.boot.test.context.SpringBootTest;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.Mockito.*;
+
+@SpringBootTest(classes = DataLoaderApp.class)
+public class DataLoaderUnitTest {
+
+    private UserService userService;
+    private UserDataLoader loaderFactory;
+
+    @BeforeEach
+    void setUp() {
+        userService = mock(UserService.class);
+        loaderFactory = new UserDataLoader(userService);
+    }
+
+    @Test
+    void givenThreeUsers_whenUsingDataLoader_thenLoadUsersInSingleDBCall() throws Exception {
+        User user1 = new User("101", "User_101");
+        User user2 = new User("102", "User_102");
+        User user3 = new User("103", "User_103");
+
+        List<User> users = Arrays.asList(user1, user2, user3);
+
+        when(userService.getUsersByIds(Arrays.asList("101", "102", "103"))).thenReturn(CompletableFuture.completedFuture(users));
+
+        DataLoader<String, User> userLoader = loaderFactory.createUserLoader();
+
+        CompletableFuture<User> future1 = userLoader.load("101");
+        CompletableFuture<User> future2 = userLoader.load("102");
+        CompletableFuture<User> future3 = userLoader.load("103");
+
+        userLoader.dispatch();
+
+        assertEquals("User_101", future1.get()
+            .getName());
+        assertEquals("User_102", future2.get()
+            .getName());
+        assertEquals("User_103", future3.get()
+            .getName());
+
+        verify(userService, times(1)).getUsersByIds(Arrays.asList("101", "102", "103"));
+    }
+}