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>${java-dataloader.version}</version>
+        </dependency>
     </dependencies>
 
     <build>
@@ -123,6 +128,7 @@
         <spring.batch.version>5.2.0</spring.batch.version>
         <awaitility.version>4.2.1</awaitility.version>
         <assertj.version>3.24.2</assertj.version>
+        <java-dataloader.version>3.2.0</java-dataloader.version>
         <start-class>com.baeldung.batch.SpringBootBatchProcessingApplication</start-class>
     </properties>
 

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,21 @@
+package com.baeldung.dataloaderbatchprocessing.entity;
+
+import jakarta.persistence.Entity;
+import jakarta.persistence.Id;
+import jakarta.persistence.Table;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+
+@Entity
+@Table(name = "users")
+@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,60 @@
+package com.baeldung.dataloaderbatchprocessing;
+
+import com.baeldung.dataloaderbatchprocessing.entity.User;
+import com.baeldung.dataloaderbatchprocessing.repository.UserRepository;
+import com.baeldung.dataloaderbatchprocessing.service.UserService;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.boot.test.mock.mockito.SpyBean;
+
+import java.util.Arrays;
+import java.util.concurrent.CompletableFuture;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.*;
+
+@SpringBootTest(classes = DataLoaderApp.class)
+class UserDataLoaderIntegrationTest {
+
+    @Autowired
+    private UserRepository userRepository;
+
+    @SpyBean
+    private UserService userService;
+
+    private DataLoader<String, User> userDataLoader;
+
+    @BeforeEach
+    void setUp() {
+        userRepository.deleteAll();
+
+        User user1 = new User("101", "User_101");
+        User user2 = new User("102", "User_102");
+        User user3 = new User("103", "User_103");
+        userRepository.saveAll(Arrays.asList(user1, user2, user3));
+
+        userDataLoader = new DataLoader<>(userService::getUsersByIds);
+        DataLoaderRegistry registry = new DataLoaderRegistry();
+        registry.register("userDataLoader", userDataLoader);
+    }
+
+    @Test
+    void whenLoadingUsers_thenBatchLoaderIsInvokedAndResultsReturned() {
+        CompletableFuture<User> userFuture1 = userDataLoader.load("101");
+        CompletableFuture<User> userFuture2 = userDataLoader.load("102");
+        CompletableFuture<User> userFuture3 = userDataLoader.load("103");
+
+        userDataLoader.dispatchAndJoin();
+
+        verify(userService, times(1)).getUsersByIds(anyList());
+
+        assertThat(userFuture1.join().getName()).isEqualTo("User_101");
+        assertThat(userFuture2.join().getName()).isEqualTo("User_102");
+        assertThat(userFuture3.join().getName()).isEqualTo("User_103");
+    }
+}
+