Skip to content

Commit

Permalink
Refactoring in DefaultBatchLoaderRegistry
Browse files Browse the repository at this point in the history
Closes gh-789
  • Loading branch information
rstoyanchev committed Sep 1, 2023
1 parent 9b61d28 commit 372feab
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 33 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
* @since 1.1
* 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;
}

Expand All @@ -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);
}
Expand Down Expand Up @@ -166,14 +166,19 @@ 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 = () ->
new DataLoaderOptions(this.options != null ?
this.options : DefaultBatchLoaderRegistry.this.defaultOptionsSupplier.get());

if (this.optionsConsumer == null) {
return optionsSupplier;
Expand All @@ -185,14 +190,6 @@ private Supplier<DataLoaderOptions> initOptionsSupplier() {
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());
}
}


Expand All @@ -206,12 +203,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;
Expand All @@ -222,9 +218,8 @@ public String getName() {
return this.name;
}

@Nullable
public DataLoaderOptions getOptions() {
return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
return this.optionsSupplier.get();
}

@Override
Expand Down Expand Up @@ -257,12 +252,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;
Expand All @@ -273,9 +267,8 @@ public String getName() {
return this.name;
}

@Nullable
public DataLoaderOptions getOptions() {
return (this.optionsSupplier != null ? this.optionsSupplier.get() : null);
return this.optionsSupplier.get();
}

@Override
Expand Down
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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;

Expand Down Expand Up @@ -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();

Expand Down

0 comments on commit 372feab

Please sign in to comment.