package io.vertx.ext.web.handler.graphql.dataloader.impl;
import io.vertx.core.Context;
import io.vertx.core.Future;
import io.vertx.core.Promise;
import io.vertx.core.impl.ContextInternal;
import io.vertx.ext.web.handler.graphql.dataloader.VertxBatchLoader;
import org.dataloader.BatchLoaderEnvironment;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletionStage;
import java.util.function.BiFunction;
import java.util.function.Function;
public class FutureBatchLoaderImpl<K, V> implements VertxBatchLoader<K, V> {
private final BiFunction<List<K>, BatchLoaderEnvironment, Future<List<V>>> batchLoader;
private final Function<BatchLoaderEnvironment, Context> contextProvider;
public FutureBatchLoaderImpl(
BiFunction<List<K>, BatchLoaderEnvironment, Future<List<V>>> batchLoader,
Function<BatchLoaderEnvironment, Context> contextProvider
) {
this.batchLoader = Objects.requireNonNull(batchLoader, "batchLoader is null");
this.contextProvider = Objects.requireNonNull(contextProvider, "contextProvider is null");
}
@Override
public CompletionStage<List<V>> load(List<K> keys, BatchLoaderEnvironment env) {
ContextInternal context = (ContextInternal) contextProvider.apply(env);
Promise<List<V>> promise;
if (context == null) {
promise = Promise.promise();
invokeBatchLoader(keys, env, promise);
} else {
promise = context.promise();
context.runOnContext(v -> invokeBatchLoader(keys, env, promise));
}
return promise.future().toCompletionStage();
}
private void invokeBatchLoader(List<K> keys, BatchLoaderEnvironment env, Promise<List<V>> promise) {
batchLoader.apply(keys, env).onComplete(promise);
}
}