A registry and a dispatcher for graphql

Author: bbakerman

build.gradle

@@ -49,6 +49,7 @@ compileJava {
 }
 
 dependencies {
+    compile "com.graphql-java:graphql-java:4.0"
     testCompile "junit:junit:$junitVersion"
     testCompile 'org.awaitility:awaitility:2.0.0'
 }

src/main/java/org/dataloader/DataLoaderRegistry.java

@@ -0,0 +1,54 @@
+package org.dataloader;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+
+/**
+ * This allows data loaders to be registered together into a single place so
+ * they can be dispatched as one.
+ */
+public class DataLoaderRegistry {
+    private final List<DataLoader<?, ?>> dataLoaders = new CopyOnWriteArrayList<>();
+
+    /**
+     * @return the currently registered data loaders
+     */
+    public List<DataLoader<?, ?>> getDataLoaders() {
+        return new ArrayList<>(dataLoaders);
+    }
+
+    /**
+     * This will register a new dataloader
+     *
+     * @param dataLoader the data loader to register
+     *
+     * @return this registry
+     */
+    public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
+        if (!dataLoaders.contains(dataLoader)) {
+            dataLoaders.add(dataLoader);
+        }
+        return this;
+    }
+
+    /**
+     * This will unregister a new dataloader
+     *
+     * @param dataLoader the data loader to unregister
+     *
+     * @return this registry
+     */
+    public DataLoaderRegistry unregister(DataLoader<?, ?> dataLoader) {
+        dataLoaders.remove(dataLoader);
+        return this;
+    }
+
+    /**
+     * This will called {@link org.dataloader.DataLoader#dispatch()} on each of the registered
+     * {@link org.dataloader.DataLoader}s
+     */
+    public void dispatchAll() {
+        dataLoaders.forEach(DataLoader::dispatch);
+    }
+}

src/main/java/org/dataloader/graphql/DataLoaderDispatcherInstrumentation.java

@@ -0,0 +1,28 @@
+package org.dataloader.graphql;
+
+import graphql.ExecutionResult;
+import graphql.execution.instrumentation.InstrumentationContext;
+import graphql.execution.instrumentation.NoOpInstrumentation;
+import graphql.execution.instrumentation.parameters.InstrumentationExecutionStrategyParameters;
+import org.dataloader.DataLoaderRegistry;
+
+import java.util.concurrent.CompletableFuture;
+
+/**
+ * This graphql {@link graphql.execution.instrumentation.Instrumentation} will dispatch
+ * all the contained {@link org.dataloader.DataLoader}s when each level of the graphql
+ * query is executed.
+ */
+public class DataLoaderDispatcherInstrumentation extends NoOpInstrumentation {
+
+    private final DataLoaderRegistry dataLoaderRegistry;
+
+    public DataLoaderDispatcherInstrumentation(DataLoaderRegistry dataLoaderRegistry) {
+        this.dataLoaderRegistry = dataLoaderRegistry;
+    }
+
+    @Override
+    public InstrumentationContext<CompletableFuture<ExecutionResult>> beginExecutionStrategy(InstrumentationExecutionStrategyParameters parameters) {
+        return (result, t) -> dataLoaderRegistry.dispatchAll();
+    }
+}

src/test/java/org/dataloader/DataLoaderRegistryTest.java

@@ -0,0 +1,41 @@
+package org.dataloader;
+
+import org.hamcrest.Matchers;
+import org.junit.Test;
+
+import java.util.concurrent.CompletableFuture;
+
+import static java.util.Arrays.asList;
+import static org.junit.Assert.assertThat;
+
+public class DataLoaderRegistryTest {
+    final BatchLoader<Object, Object> identityBatchLoader = CompletableFuture::completedFuture;
+
+    @Test
+    public void registration_works() throws Exception {
+        DataLoader<Object, Object> dlA = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> dlB = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> dlC = new DataLoader<>(identityBatchLoader);
+
+        DataLoaderRegistry registry = new DataLoaderRegistry();
+
+        registry.register(dlA).register(dlB).register(dlC);
+
+        assertThat(registry.getDataLoaders(), Matchers.equalTo(asList(dlA, dlB, dlC)));
+
+        // the same dl twice is one add
+
+
+        registry = new DataLoaderRegistry();
+
+        registry.register(dlA).register(dlB).register(dlC).register(dlA).register(dlB);
+
+        assertThat(registry.getDataLoaders(), Matchers.equalTo(asList(dlA, dlB, dlC)));
+
+
+        // and unregister
+        registry.unregister(dlC);
+
+        assertThat(registry.getDataLoaders(), Matchers.equalTo(asList(dlA, dlB)));
+    }
+}

src/test/java/org/dataloader/graphql/DataLoaderDispatcherInstrumentationTest.java

@@ -0,0 +1,56 @@
+package org.dataloader.graphql;
+
+import graphql.ExecutionResult;
+import graphql.execution.instrumentation.InstrumentationContext;
+import org.dataloader.BatchLoader;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import static java.util.Arrays.asList;
+import static java.util.Collections.singletonList;
+import static org.hamcrest.Matchers.equalTo;
+import static org.junit.Assert.assertThat;
+
+public class DataLoaderDispatcherInstrumentationTest {
+
+    @Test
+    public void basic_invocation() throws Exception {
+
+        AtomicInteger invocationCount = new AtomicInteger();
+        final List<Object> loadedKeys = new ArrayList<>();
+        final BatchLoader<Object, Object> identityBatchLoader = keys -> {
+            invocationCount.incrementAndGet();
+            loadedKeys.add(keys);
+            return CompletableFuture.completedFuture(keys);
+        };
+
+        DataLoader<Object, Object> dlA = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> dlB = new DataLoader<>(identityBatchLoader);
+        DataLoader<Object, Object> dlC = new DataLoader<>(identityBatchLoader);
+        DataLoaderRegistry registry = new DataLoaderRegistry()
+                .register(dlA)
+                .register(dlB)
+                .register(dlC);
+
+        DataLoaderDispatcherInstrumentation dispatcher = new DataLoaderDispatcherInstrumentation(registry);
+        InstrumentationContext<CompletableFuture<ExecutionResult>> context = dispatcher.beginExecutionStrategy(null);
+
+        // cause some activity
+        dlA.load("A");
+        dlB.load("B");
+        dlC.load("C");
+
+        context.onEnd(null, null);
+
+        assertThat(invocationCount.get(), equalTo(3));
+
+        // will be [[A],[B],[C]]
+        assertThat(loadedKeys, equalTo(asList(singletonList("A"), singletonList("B"), singletonList("C"))));
+    }
+}