Refactoring in DefaultBatchLoaderRegistry

Closes gh-788

Author: rstoyanchev

spring-graphql/src/main/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistry.java

@@ -43,9 +43,9 @@
 import org.springframework.util.StringUtils;
 
 /**
- * A default implementation of {@link BatchLoaderRegistry} that accepts
- * registrations, and also an implementation of {@link DataLoaderRegistrar} to
- * apply those registrations to a {@link DataLoaderRegistry}.
+ * Default implementation of {@link BatchLoaderRegistry} that stores batch loader
+ * registrations. Also, an implementation of {@link DataLoaderRegistrar} that
+ * registers the batch loaders as {@link DataLoader}s in {@link DataLoaderRegistry}.
  *
  * @author Rossen Stoyanchev
  * @since 1.0.0
@@ -60,18 +60,19 @@ public class DefaultBatchLoaderRegistry implements BatchLoaderRegistry {
 
 
 	/**
-	 * Default constructor
+	 * Default constructor.
 	 */
 	public DefaultBatchLoaderRegistry() {
 		this(DataLoaderOptions::newOptions);
 	}
 
 	/**
 	 * Constructor with a default {@link DataLoaderOptions} supplier to use as
-	 * a starting point for all registrations.
+	 * a starting point for batch loader registrations.
 	 * @since 1.1.0
 	 */
 	public DefaultBatchLoaderRegistry(Supplier<DataLoaderOptions> defaultOptionsSupplier) {
+		Assert.notNull(defaultOptionsSupplier, "'defaultOptionsSupplier' is required");
 		this.defaultOptionsSupplier = defaultOptionsSupplier;
 	}
 
@@ -89,16 +90,15 @@ public <K, V> RegistrationSpec<K, V> forName(String name) {
 	@Override
 	public void registerDataLoaders(DataLoaderRegistry registry, GraphQLContext context) {
 		BatchLoaderContextProvider contextProvider = () -> context;
-		DataLoaderOptions defaultOptions = this.defaultOptionsSupplier.get();
 		for (ReactorBatchLoader<?, ?> loader : this.loaders) {
 			DataLoaderOptions options = loader.getOptions();
-			options = (options != null ? options : defaultOptions).setBatchLoaderContextProvider(contextProvider);
+			options = options.setBatchLoaderContextProvider(contextProvider);
 			DataLoader<?, ?> dataLoader = DataLoaderFactory.newDataLoader(loader, options);
 			registerDataLoader(loader.getName(), dataLoader, registry);
 		}
 		for (ReactorMappedBatchLoader<?, ?> loader : this.mappedLoaders) {
 			DataLoaderOptions options = loader.getOptions();
-			options = (options != null ? options : defaultOptions).setBatchLoaderContextProvider(contextProvider);
+			options = options.setBatchLoaderContextProvider(contextProvider);
 			DataLoader<?, ?> dataLoader = DataLoaderFactory.newMappedDataLoader(loader, options);
 			registerDataLoader(loader.getName(), dataLoader, registry);
 		}
@@ -166,33 +166,29 @@ public void registerMappedBatchLoader(BiFunction<Set<K>, BatchLoaderEnvironment,
 					new ReactorMappedBatchLoader<>(initName(), loader, initOptionsSupplier()));
 		}
 
-		@Nullable
-		private Supplier<DataLoaderOptions> initOptionsSupplier() {
-			if (this.options == null && this.optionsConsumer == null) {
-				return null;
+		private String initName() {
+			if (StringUtils.hasText(this.name)) {
+				return this.name;
 			}
+			Assert.notNull(this.valueType, "Value type not available to select a default DataLoader name.");
+			return (StringUtils.hasText(this.name) ? this.name : this.valueType.getName());
+		}
 
-			Supplier<DataLoaderOptions> optionsSupplier =
-					(this.options != null ? () -> this.options : defaultOptionsSupplier);
+		private Supplier<DataLoaderOptions> initOptionsSupplier() {
+
+			Supplier<DataLoaderOptions> optionsSupplier = (this.options != null ?
+					() -> this.options : DefaultBatchLoaderRegistry.this.defaultOptionsSupplier);
 
 			if (this.optionsConsumer == null) {
 				return optionsSupplier;
 			}
 
 			return () -> {
-				DataLoaderOptions options = optionsSupplier.get();
+				DataLoaderOptions options = new DataLoaderOptions(optionsSupplier.get());
 				this.optionsConsumer.accept(options);
 				return options;
 			};
 		}
-
-		private String initName() {
-			if (StringUtils.hasText(this.name)) {
-				return this.name;
-			}
-			Assert.notNull(this.valueType, "Value type not available to select a default DataLoader name.");
-			return (StringUtils.hasText(this.name) ? this.name : this.valueType.getName());
-		}
 	}
 
 
@@ -206,12 +202,11 @@ private static class ReactorBatchLoader<K, V> implements BatchLoaderWithContext<
 
 		private final BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader;
 
-		@Nullable
 		private final Supplier<DataLoaderOptions> optionsSupplier;
 
 		private ReactorBatchLoader(String name,
 				BiFunction<List<K>, BatchLoaderEnvironment, Flux<V>> loader,
-				@Nullable Supplier<DataLoaderOptions> optionsSupplier) {
+				Supplier<DataLoaderOptions> optionsSupplier) {
 
 			this.name = name;
 			this.loader = loader;
@@ -222,9 +217,8 @@ public String getName() {
 			return this.name;
 		}
 
-		@Nullable
 		public DataLoaderOptions getOptions() {
-			return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
+			return this.optionsSupplier.get();
 		}
 
 		@Override
@@ -257,12 +251,11 @@ private static class ReactorMappedBatchLoader<K, V> implements MappedBatchLoader
 
 		private final BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader;
 
-		@Nullable
 		private final Supplier<DataLoaderOptions> optionsSupplier;
 
 		private ReactorMappedBatchLoader(String name,
 				BiFunction<Set<K>, BatchLoaderEnvironment, Mono<Map<K, V>>> loader,
-				@Nullable Supplier<DataLoaderOptions> optionsSupplier) {
+				Supplier<DataLoaderOptions> optionsSupplier) {
 
 			this.name = name;
 			this.loader = loader;
@@ -273,9 +266,8 @@ public String getName() {
 			return this.name;
 		}
 
-		@Nullable
 		public DataLoaderOptions getOptions() {
-			return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
+			return this.optionsSupplier.get();
 		}
 
 		@Override

spring-graphql/src/test/java/org/springframework/graphql/execution/DefaultBatchLoaderRegistryTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2022 the original author or authors.
+ * Copyright 2002-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 package org.springframework.graphql.execution;
 
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
@@ -101,7 +102,41 @@ void mappedBatchLoader() throws Exception {
 	}
 
 	@Test
-	void batchLoaderWithCustomNameAndOptions() {
+	void dataLoaderOptions() throws Exception {
+
+		DataLoaderOptions defaultOptions = DataLoaderOptions.newOptions().setBatchingEnabled(false);
+		DefaultBatchLoaderRegistry batchLoaderRegistry = new DefaultBatchLoaderRegistry(() -> defaultOptions);
+
+		AtomicInteger counter = new AtomicInteger(1);
+
+		batchLoaderRegistry.forName("loader1")
+				.withOptions(options -> options.setCachingEnabled(false))
+				.registerBatchLoader((keys, environment) -> Flux.just(counter.getAndIncrement()));
+
+		batchLoaderRegistry.forName("loader2")
+				.withOptions(options -> options.setCachingEnabled(true))
+				.registerBatchLoader((keys, environment) -> Flux.just(counter.getAndIncrement()));
+
+		GraphQLContext graphQLContext = GraphQLContext.newContext().build();
+		batchLoaderRegistry.registerDataLoaders(this.dataLoaderRegistry, graphQLContext);
+
+		DataLoader<Long, Integer> loader1 =
+				(DataLoader<Long, Integer>) this.dataLoaderRegistry.getDataLoadersMap().get("loader1");
+
+		assertThat(loader1.load(1L).get()).isEqualTo(1);
+		assertThat(loader1.load(1L).get()).isEqualTo(2);
+		assertThat(loader1.load(1L).get()).isEqualTo(3);
+
+		DataLoader<Long, Integer> loader2 =
+				(DataLoader<Long, Integer>) this.dataLoaderRegistry.getDataLoadersMap().get("loader2");
+
+		assertThat(loader2.load(1L).get()).isEqualTo(4);
+		assertThat(loader2.load(1L).get()).isEqualTo(4);
+		assertThat(loader2.load(1L).get()).isEqualTo(4);
+	}
+
+	@Test
+	void batchLoaderOptionsConsumer() {
 		String name = "myLoader";
 		StatisticsCollector collector = new NoOpStatisticsCollector();