@BatchMapping methods should pass the localContext as the BatchLoaderEnvironment's keyContexts

[ooraini]

calling getKeyContexts on BatchLoaderEnvironment returns a list of nulls. I think this is caused by Spring's BatchMappingDataFetcher not propagating the local context. Specifically line 606:

static class BatchMappingDataFetcher implements DataFetcher<Object>, SelfDescribingDataFetcher<Object> {

		private final DataFetcherMappingInfo mappingInfo;

		private final ResolvableType returnType;

		private final String dataLoaderKey;

		BatchMappingDataFetcher(DataFetcherMappingInfo info, ResolvableType valueType, String dataLoaderKey) {
			this.mappingInfo = info;
			this.returnType = ResolvableType.forClassWithGenerics(CompletableFuture.class, valueType);
			this.dataLoaderKey = dataLoaderKey;
		}

		@Override
		public String getDescription() {
			return "@BatchMapping " + this.mappingInfo.getHandlerMethod().getShortLogMessage();
		}

		@Override
		public ResolvableType getReturnType() {
			return this.returnType;
		}

		@Override
		public Object get(DataFetchingEnvironment env) {
			DataLoader<?, ?> dataLoader = env.getDataLoaderRegistry().getDataLoader(this.dataLoaderKey);
			Assert.state(dataLoader != null, "No DataLoader for key '" + this.dataLoaderKey + "'");
			return dataLoader.load(env.getSource()); <-- HERE
		}

		@Override
		public String toString() {
			return getDescription();
		}
	}

Changing it to: return dataLoader.load(env.getSource(), env.getLocalContext()) might fix the issue.(I think)

[bclozel]

Thanks for reaching out @ooraini .

I'm not sure I understand the problem here. Rather than pointing to a possible fix, could you explain what you're doing and where you're getting the BatchLoaderEnvironment from? Maybe share a sample controller that demonstrates the problem or better, a minimal sample application?

Here, this part of our codebase is about registering @BatchMapping controller methods, and I don't see the link with BatchLoaderEnvironment and what you would expect the key contexts to contain. The fix you're suggesting seems to go against the API, because I think getKeyContexts should return a Map<String, GraphQLContext>.

[ooraini]

Apologies, I meant to refer to the getKeyContextsList method. Here is an example. Suppose I have the following SchemaMapping:

    @SchemaMapping(typeName = "ShareholderRegisterShareholder")
    S3ObjectRef certificate(ShareholderProjection projection,
                            @LocalContextValue ShareholderRegisterProjection register,
                            @LocalContextValue Company company) {
         // ...more code
    }

I'm using two values from the local context. Things work. Now, I want to optimize it by processing fields in batch using BatchMapping. I was hoping that the following would work:

    @BatchMapping(typeName = "ShareholderRegisterShareholder")
    Map<LegalPersonIdentifier, S3ObjectRef> certificate(List<ShareholderProjection> projections,
                                                        BatchLoaderEnvironment environment) {
        List<Object> keyContextsList = environment.getKeyContextsList();
        for (int i = 0, projectionsSize = projections.size(); i < projectionsSize; i++) {
            ShareholderProjection projection = projections.get(i);
            GraphQLContext graphQLContext = (GraphQLContext) keyContextsList.get(i);
            Company company = graphQLContext.get("company");
            ShareholderRegisterProjection register = graphQLContext.get("register");

            // do stuff, e.g. group the data and batch process it

        }
    }

But, the getKeyContextsList is returning a list of null elements. I traced back the execution into graphql-java and the datalaoder library until I got to BatchMappingDataFetcher.

what you would expect the key contexts to contain

I expect(hope?) getKeyContextsList to contain the local context for each source object in List<SOURCE>.

Hope that helps.