Small cleanup with prev PR

server/implementation/src/main/java/io/smallrye/graphql/bootstrap/DataFetcherFactory.java

@@ -164,14 +164,6 @@ public PlugableDataFetcher getOtherFieldDataFetcher(Operation operation) {
     }
 
     public PlugableDataFetcher getDefaultDataFetcher(Operation operation) {
-
-        //        for (DataFetcherService dfe : dataFetcherServices) {
-        //            PlugableDataFetcher df = dfe.getOtherFieldDataFetcher(operation);
-        //            if (df != null) {
-        //                return df;
-        //            }
-        //        }
-
         return getOtherFieldDataFetcher(operation);
     }
 

server/implementation/src/main/java/io/smallrye/graphql/execution/ExecutionService.java

@@ -32,6 +32,7 @@
 import io.smallrye.graphql.execution.context.SmallRyeBatchLoaderContextProvider;
 import io.smallrye.graphql.execution.context.SmallRyeContext;
 import io.smallrye.graphql.execution.datafetcher.helper.BatchLoaderHelper;
+import io.smallrye.graphql.execution.datafetcher.helper.ContextHelper;
 import io.smallrye.graphql.execution.error.ExceptionHandler;
 import io.smallrye.graphql.execution.event.EventEmitter;
 import io.smallrye.graphql.schema.model.Operation;
@@ -170,7 +171,7 @@ public void execute(JsonObject jsonInput, Map<String, Object> context, Execution
                 smallRyeContext = smallRyeContext.withDataFromExecution(executionInput, queryCache);
 
                 // Context
-                context.put("context", smallRyeContext);
+                context.put(ContextHelper.CONTEXT, smallRyeContext);
                 executionInput.getGraphQLContext().putAll(context);
 
                 // Notify before
@@ -204,7 +205,7 @@ private void writeAsync(GraphQL graphQL,
         executionResult.whenComplete((t, u) -> {
             executionInput.getGraphQLContext().putAll(context);
 
-            SmallRyeContext smallryeContext = (SmallRyeContext) context.get("context");
+            SmallRyeContext smallryeContext = (SmallRyeContext) context.get(ContextHelper.CONTEXT);
             SmallRyeContext.setContext(smallryeContext);
 
             // Notify after
@@ -229,7 +230,7 @@ private void writeSync(GraphQL g,
         try {
             ExecutionResult executionResult = g.execute(executionInput);
             // Notify after
-            eventEmitter.fireAfterExecute((Context) context.get("context"));
+            eventEmitter.fireAfterExecute((Context) context.get(ContextHelper.CONTEXT));
 
             ExecutionResponse executionResponse = new ExecutionResponse(executionResult);
             if (!payloadOption.equals(LogPayloadOption.off)) {

server/implementation/src/main/java/io/smallrye/graphql/execution/context/SmallRyeBatchLoaderContextProvider.java

@@ -6,7 +6,7 @@
 import org.dataloader.BatchLoaderContextProvider;
 import org.dataloader.DataLoader;
 
-import graphql.GraphQLContext;
+import graphql.schema.DataFetchingEnvironment;
 
 // hacky way to pass the GraphQL Context from data from data fetchers to BatchLoaderEnvironment
 public class SmallRyeBatchLoaderContextProvider implements BatchLoaderContextProvider {
@@ -28,14 +28,14 @@ public static SmallRyeBatchLoaderContextProvider getForDataLoader(DataLoader dat
         return INSTANCES.get(dataLoader);
     }
 
-    private AtomicReference<GraphQLContext> current = new AtomicReference<>();
+    private AtomicReference<DataFetchingEnvironment> current = new AtomicReference<>();
 
-    public void set(GraphQLContext context) {
-        current.set(context);
+    public void set(DataFetchingEnvironment dfe) {
+        current.set(dfe);
     }
 
     @Override
-    public GraphQLContext getContext() {
+    public DataFetchingEnvironment getContext() {
         return current.getAndSet(null);
     }
 }

server/implementation/src/main/java/io/smallrye/graphql/execution/context/SmallRyeContext.java

@@ -100,10 +100,6 @@ public SmallRyeContext withDataFromFetcher(DataFetchingEnvironment dfe, Field fi
         return new SmallRyeContext(this.jsonObject, dfe, this.executionInput, this.queryCache, field);
     }
 
-    public SmallRyeContext withDataFromFetcher(Field field) {
-        return withDataFromFetcher(this.dfe, field);
-    }
-
     public static void remove() {
         current.remove();
     }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/AbstractDataFetcher.java

@@ -7,6 +7,7 @@
 import io.smallrye.graphql.execution.context.SmallRyeContext;
 import io.smallrye.graphql.execution.datafetcher.helper.ArgumentHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.BatchLoaderHelper;
+import io.smallrye.graphql.execution.datafetcher.helper.ContextHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.ErrorResultHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.FieldHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.OperationInvoker;
@@ -21,7 +22,7 @@
  * @param <K>
  * @param <T>
  */
-public abstract class AbstractDataFetcher<K, T> implements PlugableDataFetcher<K, T>, ContextAware {
+public abstract class AbstractDataFetcher<K, T> implements PlugableDataFetcher<K, T> {
 
     protected Operation operation;
     protected FieldHelper fieldHelper;
@@ -30,6 +31,7 @@ public abstract class AbstractDataFetcher<K, T> implements PlugableDataFetcher<K
     protected ArgumentHelper argumentHelper;
     protected EventEmitter eventEmitter = EventEmitter.getInstance();
     protected BatchLoaderHelper batchLoaderHelper = new BatchLoaderHelper();
+    protected ContextHelper contextHelper = new ContextHelper();
 
     public AbstractDataFetcher(Operation operation) {
         this.operation = operation;
@@ -68,7 +70,7 @@ public T get(final DataFetchingEnvironment dfe) throws Exception {
 
     private SmallRyeContext initSmallRyeContext(final DataFetchingEnvironment dfe) {
         // update the context
-        SmallRyeContext context = updateSmallRyeContext(dfe, operation);
+        SmallRyeContext context = contextHelper.updateSmallRyeContextWithField(dfe, operation);
         eventEmitter.fireBeforeDataFetch(context);
         return context;
     }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/BatchDataFetcher.java

@@ -8,6 +8,7 @@
 import io.smallrye.graphql.execution.context.SmallRyeContext;
 import io.smallrye.graphql.execution.datafetcher.helper.ArgumentHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.BatchLoaderHelper;
+import io.smallrye.graphql.execution.datafetcher.helper.ContextHelper;
 import io.smallrye.graphql.execution.event.EventEmitter;
 import io.smallrye.graphql.schema.model.Operation;
 
@@ -16,12 +17,13 @@
  * 
  * @author Phillip Kruger ([email protected])
  */
-public class BatchDataFetcher<T> implements DataFetcher<T>, ContextAware {
+public class BatchDataFetcher<T> implements DataFetcher<T> {
 
     private final Operation operation;
     private final ArgumentHelper argumentHelper;
     private final String batchLoaderName;
     private final BatchLoaderHelper batchLoaderHelper = new BatchLoaderHelper();
+    private final ContextHelper contextHelper = new ContextHelper();
     private final EventEmitter eventEmitter = EventEmitter.getInstance();
 
     public BatchDataFetcher(Operation operation) {
@@ -33,13 +35,13 @@ public BatchDataFetcher(Operation operation) {
     @Override
     public T get(final DataFetchingEnvironment dfe) throws Exception {
 
-        SmallRyeContext smallryeContext = updateSmallRyeContext(dfe, operation);
+        SmallRyeContext smallryeContext = contextHelper.updateSmallRyeContextWithField(dfe, operation);
         eventEmitter.fireBeforeDataFetch(smallryeContext);
         Object[] transformedArguments = argumentHelper.getArguments(dfe, true);
         Object source = dfe.getSource();
 
         DataLoader<Object, Object> dataLoader = dfe.getDataLoader(batchLoaderName);
-        updateSmallRyeContext(dataLoader, dfe);
+        batchLoaderHelper.setDataFetchingEnvironment(dataLoader, dfe);
 
         try {
             SmallRyeContext.setContext(smallryeContext);

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/CompletionStageDataFetcher.java

@@ -32,7 +32,7 @@ public CompletionStageDataFetcher(Operation operation) {
     protected <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResult.Builder<Object> resultBuilder,
             Object[] transformedArguments) throws AbstractDataFetcherException, Exception {
 
-        SmallRyeContext context = getSmallRyeContext(dfe);
+        SmallRyeContext context = contextHelper.getSmallRyeContext(dfe);
         ThreadContext threadContext = ThreadContext.builder().build();
         SmallRyeContext.setContext(context);
         try {
@@ -75,7 +75,7 @@ protected <T> T invokeFailure(DataFetcherResult.Builder<Object> resultBuilder) {
     public CompletionStage<List<T>> load(List<K> keys, BatchLoaderEnvironment ble) {
         Object[] arguments = batchLoaderHelper.getArguments(keys, ble);
         final ClassLoader tccl = Thread.currentThread().getContextClassLoader();
-        final SmallRyeContext smallRyeContext = getSmallRyeContext(ble);
+        final SmallRyeContext smallRyeContext = contextHelper.getSmallRyeContext(ble);
 
         final ThreadContext threadContext = ThreadContext.builder().build();
         try {

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/DefaultDataFetcher.java

@@ -30,7 +30,7 @@ public DefaultDataFetcher(Operation operation) {
     @Override
     public <T> T invokeAndTransform(DataFetchingEnvironment dfe, DataFetcherResult.Builder<Object> resultBuilder,
             Object[] transformedArguments) throws Exception {
-        SmallRyeContext context = getSmallRyeContext(dfe);
+        SmallRyeContext context = contextHelper.getSmallRyeContext(dfe);
         try {
             SmallRyeContext.setContext(context);
             Object resultFromMethodCall = operationInvoker.invoke(transformedArguments);
@@ -49,9 +49,10 @@ public <T> T invokeFailure(DataFetcherResult.Builder<Object> resultBuilder) {
 
     @Override
     public CompletionStage<List<T>> load(List<K> keys, BatchLoaderEnvironment ble) {
-        Object[] arguments = batchLoaderHelper.getArguments(keys, ble);
+        final Object[] arguments = batchLoaderHelper.getArguments(keys, ble);
+        final DataFetchingEnvironment dfe = batchLoaderHelper.getDataFetchingEnvironment(ble);
+        final SmallRyeContext smallRyeContext = contextHelper.getSmallRyeContext(dfe);
         final ClassLoader tccl = Thread.currentThread().getContextClassLoader();
-        final SmallRyeContext smallRyeContext = getSmallRyeContext(ble);
 
         ThreadContext threadContext = ThreadContext.builder().build();
         try {

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/FieldDataFetcher.java

@@ -8,6 +8,7 @@
 import graphql.TrivialDataFetcher;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import io.smallrye.graphql.execution.datafetcher.helper.ContextHelper;
 import io.smallrye.graphql.execution.datafetcher.helper.FieldHelper;
 import io.smallrye.graphql.schema.model.Field;
 import io.smallrye.graphql.schema.model.Reference;
@@ -26,10 +27,10 @@
  *           different
  *           subtype of the owner class for each call).
  */
-public class FieldDataFetcher<T> implements DataFetcher<T>, TrivialDataFetcher<T>, ContextAware {
+public class FieldDataFetcher<T> implements DataFetcher<T>, TrivialDataFetcher<T> {
 
     private final FieldHelper fieldHelper;
-
+    private final ContextHelper contextHelper = new ContextHelper();
     private final Field field;
 
     /**
@@ -51,7 +52,7 @@ public FieldDataFetcher(final Field field, final Reference owner) {
     @Override
     public T get(DataFetchingEnvironment dfe) throws Exception {
 
-        updateSmallRyeContext(dfe, field);
+        contextHelper.updateSmallRyeContextWithField(dfe, field);
 
         if (this.propertyAccessor == null) {
             // lazy initialize method handle, does not have to be threadsafe

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/MultiDataFetcher.java

@@ -35,7 +35,7 @@ protected <O> O invokeAndTransform(
             DataFetchingEnvironment dfe,
             DataFetcherResult.Builder<Object> resultBuilder,
             Object[] transformedArguments) throws Exception {
-        SmallRyeContext context = getSmallRyeContext(dfe);
+        SmallRyeContext context = contextHelper.getSmallRyeContext(dfe);
         try {
             SmallRyeContext.setContext(context);
             Multi<?> multi = operationInvoker.invoke(transformedArguments);

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/PublisherDataFetcher.java

@@ -37,7 +37,7 @@ protected <O> O invokeAndTransform(
             DataFetchingEnvironment dfe,
             DataFetcherResult.Builder<Object> resultBuilder,
             Object[] transformedArguments) throws Exception {
-        SmallRyeContext context = getSmallRyeContext(dfe);
+        SmallRyeContext context = contextHelper.getSmallRyeContext(dfe);
         try {
             SmallRyeContext.setContext(context);
             Publisher<?> publisher = operationInvoker.invoke(transformedArguments);

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/UniDataFetcher.java

@@ -34,7 +34,7 @@ protected <O> O invokeAndTransform(
             DataFetchingEnvironment dfe,
             DataFetcherResult.Builder<Object> resultBuilder,
             Object[] transformedArguments) throws Exception {
-        SmallRyeContext context = getSmallRyeContext(dfe);
+        SmallRyeContext context = contextHelper.getSmallRyeContext(dfe);
         try {
             SmallRyeContext.setContext(context);
             Uni<?> uni = operationInvoker.invoke(transformedArguments);

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/helper/BatchLoaderHelper.java

@@ -5,7 +5,10 @@
 import java.util.List;
 
 import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.DataLoader;
 
+import graphql.schema.DataFetchingEnvironment;
+import io.smallrye.graphql.execution.context.SmallRyeBatchLoaderContextProvider;
 import io.smallrye.graphql.schema.model.Operation;
 
 /**
@@ -32,4 +35,15 @@ public <K> Object[] getArguments(List<K> keys, BatchLoaderEnvironment ble) {
 
         return arguments.toArray();
     }
+
+    public DataFetchingEnvironment getDataFetchingEnvironment(BatchLoaderEnvironment ble) {
+        return ble.getContext();
+    }
+
+    public void setDataFetchingEnvironment(final DataLoader<Object, Object> dataLoader, final DataFetchingEnvironment dfe) {
+        // FIXME: this is potentially brittle because it assumes that the batch loader will execute and
+        //  consume the context before we call this again for a different operation, but I don't know
+        //  how else to pass this context to the matching BatchLoaderEnvironment instance
+        SmallRyeBatchLoaderContextProvider.getForDataLoader(dataLoader).set(dfe);
+    }
 }

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/helper/ContextHelper.java

@@ -0,0 +1,45 @@
+package io.smallrye.graphql.execution.datafetcher.helper;
+
+import org.dataloader.BatchLoaderEnvironment;
+
+import graphql.GraphQLContext;
+import graphql.schema.DataFetchingEnvironment;
+import io.smallrye.graphql.execution.context.SmallRyeContext;
+import io.smallrye.graphql.schema.model.Field;
+
+/**
+ * Helping with context
+ * 
+ * @author Phillip Kruger ([email protected])
+ */
+public class ContextHelper {
+
+    public static final String CONTEXT = "context";
+
+    public SmallRyeContext getSmallRyeContext(final BatchLoaderEnvironment ble) {
+        DataFetchingEnvironment dfe = ble.getContext();
+        return getSmallRyeContext(dfe);
+    }
+
+    public SmallRyeContext getSmallRyeContext(final DataFetchingEnvironment dfe) {
+        GraphQLContext graphQLContext = dfe.getGraphQlContext();
+        return graphQLContext.get(CONTEXT);
+    }
+
+    public void setSmallRyeContext(final DataFetchingEnvironment dfe, SmallRyeContext smallRyeContext) {
+        GraphQLContext graphQLContext = dfe.getGraphQlContext();
+        graphQLContext.put(CONTEXT, smallRyeContext);
+    }
+
+    public SmallRyeContext updateSmallRyeContextWithField(final DataFetchingEnvironment dfe, final Field field) {
+        SmallRyeContext smallRyeContext = getSmallRyeContext(dfe);
+        smallRyeContext = smallRyeContext.withDataFromFetcher(dfe, field);
+        setSmallRyeContext(dfe, smallRyeContext);
+        return smallRyeContext;
+    }
+
+    public SmallRyeContext updateSmallRyeContextWithField(final BatchLoaderEnvironment ble, final Field field) {
+        DataFetchingEnvironment dfe = ble.getContext();
+        return updateSmallRyeContextWithField(dfe, field);
+    }
+}

server/implementation/src/main/java/io/smallrye/graphql/execution/datafetcher/helper/ReflectionInvoker.java

@@ -81,6 +81,9 @@ public Object run() throws Exception {
     }
 
     public <T> T invoke(Object... arguments) throws Exception {
+
+        // TODO: Context propagation ?
+
         if (this.injectContextAt > -1) {
             arguments = injectContext(arguments);
         }