You can now change the context during ExecutionInput build but the values are copied into the known instance

bbakerman · bbakerman · commit b471852d09ec · 2021-06-11T12:19:06.000+10:00
diff --git a/src/main/java/graphql/ExecutionInput.java b/src/main/java/graphql/ExecutionInput.java
@@ -153,7 +153,7 @@ public ExecutionInput transform(Consumer<Builder> builderConsumer) {
                 .query(this.query)
                 .operationName(this.operationName)
                 .context(this.context)
-                .graphQLContext(this.graphQLContext)
+                .transfer(this.graphQLContext)
                 .localContext(this.localContext)
                 .root(this.root)
                 .dataLoaderRegistry(this.dataLoaderRegistry)
@@ -282,12 +282,6 @@ public Builder context(Object context) {
             return this;
         }
 
-        // hidden on purpose
-        private Builder graphQLContext(GraphQLContext graphQLContext) {
-            this.graphQLContext = graphQLContext;
-            return this;
-        }
-
         /**
          * The legacy context object
          *
@@ -319,6 +313,39 @@ public Builder context(UnaryOperator<GraphQLContext.Builder> contextBuilderFunct
             return context(builder.build());
         }
 
+        /**
+         * This will give you a builder of {@link GraphQLContext} and any values you set will be copied
+         * into the underlying {@link GraphQLContext} of this execution input
+         *
+         * @param builderFunction a builder function you can use to put values into the context
+         *
+         * @return this builder
+         */
+        public Builder graphQLContext(Consumer<GraphQLContext.Builder> builderFunction) {
+            GraphQLContext.Builder builder = GraphQLContext.newContext();
+            builderFunction.accept(builder);
+            this.graphQLContext.putAll(builder);
+            return this;
+        }
+
+        /**
+         * This will put all the valuyes from the map into the underlying {@link GraphQLContext} of this execution input
+         *
+         * @param mapOfContext a map of values to put in the context
+         *
+         * @return this builder
+         */
+        public Builder graphQLContext(Map<?, Object> mapOfContext) {
+            this.graphQLContext.putAll(mapOfContext);
+            return this;
+        }
+
+        // hidden on purpose
+        private Builder transfer(GraphQLContext graphQLContext) {
+            this.graphQLContext = Assert.assertNotNull(graphQLContext);
+            return this;
+        }
+
         public Builder root(Object root) {
             this.root = root;
             return this;
diff --git a/src/main/java/graphql/GraphQLContext.java b/src/main/java/graphql/GraphQLContext.java
@@ -236,6 +236,14 @@ public static Builder newContext() {
     public static class Builder {
         private final ConcurrentMap<Object, Object> map = new ConcurrentHashMap<>();
 
+        public Builder put(
+                Object key1, Object value1
+        ) {
+            return putImpl(
+                    key1, value1
+            );
+        }
+
         public Builder of(
                 Object key1, Object value1
         ) {
@@ -311,6 +319,17 @@ public Builder of(Map<?, Object> mapOfContext) {
             return this;
         }
 
+        /**
+         * Adds all of the values in the map into the context builder.  All keys and values MUST be non null
+         *
+         * @param mapOfContext the map to put into context
+         *
+         * @return this builder
+         */
+        public Builder putAll(Map<?, Object> mapOfContext) {
+            return of(mapOfContext);
+        }
+
         /**
          * Adds all of the values in the map into the context builder.  All keys and values MUST be non null
          *
diff --git a/src/test/groovy/graphql/ExecutionInputTest.groovy b/src/test/groovy/graphql/ExecutionInputTest.groovy
@@ -26,12 +26,13 @@ class ExecutionInputTest extends Specification {
                 .variables(variables)
                 .root(root)
                 .context(context)
+                .graphQLContext({ it.of(["a": "b"]) })
                 .locale(Locale.GERMAN)
                 .extensions([some: "map"])
                 .build()
         then:
         executionInput.context == context
-        executionInput.graphQLContext != null
+        executionInput.graphQLContext.get("a") == "b"
         executionInput.root == root
         executionInput.variables == variables
         executionInput.dataLoaderRegistry == registry
@@ -41,6 +42,15 @@ class ExecutionInputTest extends Specification {
         executionInput.extensions == [some: "map"]
     }
 
+    def "map context build works"() {
+        when:
+        def executionInput = ExecutionInput.newExecutionInput().query(query)
+                .graphQLContext([a: "b"])
+                .build()
+        then:
+        executionInput.graphQLContext.get("a") == "b"
+    }
+
     def "legacy context methods work"() {
         when:
         def executionInput = ExecutionInput.newExecutionInput().query(query)
@@ -83,6 +93,7 @@ class ExecutionInputTest extends Specification {
                 .extensions([some: "map"])
                 .root(root)
                 .context(context)
+                .graphQLContext({ it.of(["a": "b"]) })
                 .locale(Locale.GERMAN)
                 .build()
         def graphQLContext = executionInputOld.getGraphQLContext()
