Merge pull request #4003 from graphql-java/chained-dataloader-optimize

andimarek · web-flow · commit 9c358d367dc6 · 2025-06-05T19:24:48.000+10:00
minimize work done when chained dataloaders are disabled
diff --git a/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java b/src/main/java/graphql/schema/DataFetchingEnvironmentImpl.java
@@ -13,6 +13,7 @@
 import graphql.execution.MergedField;
 import graphql.execution.directives.QueryDirectives;
 import graphql.execution.incremental.DeferredCallContext;
+import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys;
 import graphql.language.Document;
 import graphql.language.Field;
 import graphql.language.FragmentDefinition;
@@ -217,6 +218,9 @@ public ExecutionStepInfo getExecutionStepInfo() {
 
     @Override
     public <K, V> @Nullable DataLoader<K, V> getDataLoader(String dataLoaderName) {
+        if (!graphQLContext.getBoolean(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING, false)) {
+            return dataLoaderRegistry.getDataLoader(dataLoaderName);
+        }
         return new DataLoaderWithContext<>(this, dataLoaderName, dataLoaderRegistry.getDataLoader(dataLoaderName));
     }
 
diff --git a/src/test/groovy/graphql/schema/DataFetchingEnvironmentImplTest.groovy b/src/test/groovy/graphql/schema/DataFetchingEnvironmentImplTest.groovy
@@ -4,6 +4,7 @@ import graphql.GraphQLContext
 import graphql.execution.CoercedVariables
 import graphql.execution.ExecutionId
 import graphql.execution.ExecutionStepInfo
+import graphql.execution.instrumentation.dataloader.DataLoaderDispatchingContextKeys
 import graphql.language.Argument
 import graphql.language.Field
 import graphql.language.FragmentDefinition
@@ -37,7 +38,7 @@ class DataFetchingEnvironmentImplTest extends Specification {
 
     def executionContext = newExecutionContextBuilder()
             .root("root")
-            .graphQLContext(GraphQLContext.of(["key":"context"]))
+            .graphQLContext(GraphQLContext.of(["key": "context"]))
             .executionId(executionId)
             .operationDefinition(operationDefinition)
             .document(document)
@@ -65,6 +66,7 @@ class DataFetchingEnvironmentImplTest extends Specification {
         when:
         def dfe = newDataFetchingEnvironment(executionContext)
                 .build()
+        dfe.getGraphQlContext().put(DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING, chainedDataLoaderEnabled)
         then:
         dfe.getRoot() == "root"
         dfe.getGraphQlContext().get("key") == "context"
@@ -74,13 +76,16 @@ class DataFetchingEnvironmentImplTest extends Specification {
         dfe.getOperationDefinition() == operationDefinition
         dfe.getExecutionId() == executionId
         dfe.getDataLoaderRegistry() == executionContext.getDataLoaderRegistry()
-        dfe.getDataLoader("dataLoader").delegate == executionContext.getDataLoaderRegistry().getDataLoader("dataLoader")
+        dfe.getDataLoader("dataLoader") == executionContext.getDataLoaderRegistry().getDataLoader("dataLoader") ||
+                dfe.getDataLoader("dataLoader").delegate == executionContext.getDataLoaderRegistry().getDataLoader("dataLoader")
+        where:
+        chainedDataLoaderEnabled << [true, false]
     }
 
     def "create environment from existing one will copy everything to new instance"() {
         def dfe = newDataFetchingEnvironment()
                 .context("Test Context") // Retain deprecated builder for coverage
-                .graphQLContext(GraphQLContext.of(["key": "context"]))
+                .graphQLContext(GraphQLContext.of(["key": "context", (DataLoaderDispatchingContextKeys.ENABLE_DATA_LOADER_CHAINING): chainedDataLoaderEnabled]))
                 .source("Test Source")
                 .root("Test Root")
                 .fieldDefinition(Mock(GraphQLFieldDefinition))
@@ -119,9 +124,13 @@ class DataFetchingEnvironmentImplTest extends Specification {
         dfe.getDocument() == dfeCopy.getDocument()
         dfe.getOperationDefinition() == dfeCopy.getOperationDefinition()
         dfe.getVariables() == dfeCopy.getVariables()
-        dfe.getDataLoader("dataLoader").delegate == dfeCopy.getDataLoader("dataLoader").delegate
+        dfe.getDataLoader("dataLoader") == executionContext.getDataLoaderRegistry().getDataLoader("dataLoader") ||
+                dfe.getDataLoader("dataLoader").delegate == dfeCopy.getDataLoader("dataLoader").delegate
         dfe.getLocale() == dfeCopy.getLocale()
         dfe.getLocalContext() == dfeCopy.getLocalContext()
+        where:
+        chainedDataLoaderEnabled << [true, false]
+
     }
 
     def "get or default support"() {
