package io.vertx.ext.web.handler.graphql.dataloader.impl;
import io.vertx.core.Context;
import io.vertx.core.Promise;
import io.vertx.core.impl.ContextInternal;
import io.vertx.ext.web.handler.graphql.TriConsumer;
import io.vertx.ext.web.handler.graphql.dataloader.VertxMappedBatchLoader;
import org.dataloader.BatchLoaderEnvironment;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletionStage;
import java.util.function.Function;
public class CallbackMappedBatchLoaderImpl<K, V> implements VertxMappedBatchLoader<K, V> {
private final TriConsumer<Set<K>, BatchLoaderEnvironment, Promise<Map<K, V>>> batchLoader;
private final Function<BatchLoaderEnvironment, Context> contextProvider;
public CallbackMappedBatchLoaderImpl(
TriConsumer<Set<K>, BatchLoaderEnvironment, Promise<Map<K, 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<Map<K, V>> load(Set<K> keys, BatchLoaderEnvironment env) {
ContextInternal context = (ContextInternal) contextProvider.apply(env);
Promise<Map<K, 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(Set<K> keys, BatchLoaderEnvironment env, Promise<Map<K, V>> promise) {
batchLoader.accept(keys, env, promise);
}
}