From 3f676acfa7251988fb9169695fce3ebf7d47dec2 Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 4 Jun 2025 13:57:51 +1000
Subject: [PATCH 01/15] Error Prone / NullAway support for JSpecify

---
 build.gradle                                  | 29 ++++++++++++++++++-
 .../dataloader/BatchLoaderEnvironment.java    |  7 ++---
 .../org/dataloader/DataLoaderRegistry.java    | 10 ++++---
 .../dataloader/DataLoaderRegistryTest.java    |  2 +-
 4 files changed, 38 insertions(+), 10 deletions(-)

diff --git a/build.gradle b/build.gradle
index eb57158..e07821f 100644
--- a/build.gradle
+++ b/build.gradle
@@ -11,11 +11,12 @@ plugins {
     id 'io.github.gradle-nexus.publish-plugin' version '1.0.0'
     id 'com.github.ben-manes.versions' version '0.51.0'
     id "me.champeau.jmh" version "0.7.3"
+    id "net.ltgt.errorprone" version '4.2.0'
 }
 
 java {
     toolchain {
-        languageVersion = JavaLanguageVersion.of(11)
+        languageVersion = JavaLanguageVersion.of(17)
     }
 }
 
@@ -75,8 +76,34 @@ dependencies {
     // this is needed for the idea jmh plugin to work correctly
     jmh 'org.openjdk.jmh:jmh-core:1.37'
     jmh 'org.openjdk.jmh:jmh-generator-annprocess:1.37'
+
+    errorprone 'com.uber.nullaway:nullaway:0.12.6'
+    errorprone 'com.google.errorprone:error_prone_core:2.37.0'
+}
+
+import net.ltgt.gradle.errorprone.CheckSeverity
+
+tasks.withType(JavaCompile) {
+    options.release = 11
+    options.errorprone {
+        disableAllChecks = true
+        check("NullAway", CheckSeverity.ERROR)
+        //
+        // end state has us with this config turned on - eg all classes
+        //
+        //option("NullAway:AnnotatedPackages", "org.dataloader")
+        option("NullAway:OnlyNullMarked", "true")
+        option("NullAway:JSpecifyMode", "true")
+    }
+    // Include to disable NullAway on test code
+    if (name.toLowerCase().contains("test")) {
+        options.errorprone {
+            disable("NullAway")
+        }
+    }
 }
 
+
 task sourcesJar(type: Jar) {
     dependsOn classes
     archiveClassifier.set('sources')
diff --git a/src/main/java/org/dataloader/BatchLoaderEnvironment.java b/src/main/java/org/dataloader/BatchLoaderEnvironment.java
index 6b84e70..c7a2ed8 100644
--- a/src/main/java/org/dataloader/BatchLoaderEnvironment.java
+++ b/src/main/java/org/dataloader/BatchLoaderEnvironment.java
@@ -19,11 +19,11 @@
 @NullMarked
 public class BatchLoaderEnvironment {
 
-    private final Object context;
+    private final @Nullable Object context;
     private final Map<Object, Object> keyContexts;
     private final List<Object> keyContextsList;
 
-    private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
+    private BatchLoaderEnvironment(@Nullable Object context, List<Object> keyContextsList, Map<Object, Object> keyContexts) {
         this.context = context;
         this.keyContexts = keyContexts;
         this.keyContextsList = keyContextsList;
@@ -33,7 +33,6 @@ private BatchLoaderEnvironment(Object context, List<Object> keyContextsList, Map
      * Returns the overall context object provided by {@link org.dataloader.BatchLoaderContextProvider}
      *
      * @param <T> the type you would like the object to be
-     *
      * @return a context object or null if there isn't one
      */
     @SuppressWarnings("unchecked")
@@ -68,7 +67,7 @@ public static Builder newBatchLoaderEnvironment() {
     }
 
     public static class Builder {
-        private Object context;
+        private @Nullable Object context;
         private Map<Object, Object> keyContexts = Collections.emptyMap();
         private List<Object> keyContextsList = Collections.emptyList();
 
diff --git a/src/main/java/org/dataloader/DataLoaderRegistry.java b/src/main/java/org/dataloader/DataLoaderRegistry.java
index 0988697..3b7409c 100644
--- a/src/main/java/org/dataloader/DataLoaderRegistry.java
+++ b/src/main/java/org/dataloader/DataLoaderRegistry.java
@@ -1,6 +1,7 @@
 package org.dataloader;
 
 import org.dataloader.annotations.PublicApi;
+import org.dataloader.impl.Assertions;
 import org.dataloader.instrumentation.ChainedDataLoaderInstrumentation;
 import org.dataloader.instrumentation.DataLoaderInstrumentation;
 import org.dataloader.instrumentation.DataLoaderInstrumentationHelper;
@@ -14,6 +15,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
@@ -142,7 +144,7 @@ private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options,
      */
     public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
         String name = dataLoader.getName();
-        assertState(name != null, () -> "The DataLoader must have a non null name");
+        Objects.requireNonNull(name, "The DataLoader must have a non null name");
         dataLoaders.put(name, nameAndInstrumentDL(name, instrumentation, dataLoader));
         return this;
     }
@@ -176,7 +178,7 @@ public DataLoaderRegistry register(String key, DataLoader<?, ?> dataLoader) {
      */
     public <K, V> DataLoader<K, V> registerAndGet(String key, DataLoader<?, ?> dataLoader) {
         dataLoaders.put(key, nameAndInstrumentDL(key, instrumentation, dataLoader));
-        return getDataLoader(key);
+        return Objects.requireNonNull(getDataLoader(key));
     }
 
     /**
@@ -251,10 +253,10 @@ public DataLoaderRegistry unregister(String key) {
      * @param key the key of the data loader
      * @param <K> the type of keys
      * @param <V> the type of values
-     * @return a data loader or null if its not present
+     * @return a data loader or null if it's not present
      */
     @SuppressWarnings("unchecked")
-    public <K, V> DataLoader<K, V> getDataLoader(String key) {
+    public <K, V> @Nullable DataLoader<K, V> getDataLoader(String key) {
         return (DataLoader<K, V>) dataLoaders.get(key);
     }
 
diff --git a/src/test/java/org/dataloader/DataLoaderRegistryTest.java b/src/test/java/org/dataloader/DataLoaderRegistryTest.java
index 270bd50..4335375 100644
--- a/src/test/java/org/dataloader/DataLoaderRegistryTest.java
+++ b/src/test/java/org/dataloader/DataLoaderRegistryTest.java
@@ -63,7 +63,7 @@ public void registration_works() {
         try {
             registry.register(dlUnnamed);
             Assertions.fail("Should have thrown an exception");
-        } catch (DataLoaderAssertionException ignored) {
+        } catch (NullPointerException ignored) {
         }
     }
 

From 4fbeccdc96de752a02879bbb74a974973b1d7d5b Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 4 Jun 2025 15:17:21 +1000
Subject: [PATCH 02/15] Error Prone / NullAway support for JSpecify - added
 Kotlin

---
 build.gradle                                  | 24 +++++++++++-
 src/main/java/org/dataloader/DataLoader.java  |  2 +-
 .../kotlin/org/dataloader/KotlinExamples.kt   | 39 +++++++++++++++++++
 3 files changed, 62 insertions(+), 3 deletions(-)
 create mode 100644 src/test/kotlin/org/dataloader/KotlinExamples.kt

diff --git a/build.gradle b/build.gradle
index e07821f..2d22afd 100644
--- a/build.gradle
+++ b/build.gradle
@@ -1,3 +1,6 @@
+import org.jetbrains.kotlin.gradle.dsl.JvmTarget
+import org.jetbrains.kotlin.gradle.dsl.KotlinVersion
+import net.ltgt.gradle.errorprone.CheckSeverity
 import java.text.SimpleDateFormat
 
 plugins {
@@ -12,6 +15,9 @@ plugins {
     id 'com.github.ben-manes.versions' version '0.51.0'
     id "me.champeau.jmh" version "0.7.3"
     id "net.ltgt.errorprone" version '4.2.0'
+
+    // Kotlin just for tests - not
+    id 'org.jetbrains.kotlin.jvm' version '2.1.21'
 }
 
 java {
@@ -20,6 +26,19 @@ java {
     }
 }
 
+kotlin {
+    compilerOptions {
+        apiVersion = KotlinVersion.KOTLIN_2_0
+        languageVersion = KotlinVersion.KOTLIN_2_0
+        jvmTarget = JvmTarget.JVM_11
+        javaParameters = true
+        freeCompilerArgs = [
+                '-Xemit-jvm-type-annotations',
+                '-Xjspecify-annotations=strict',
+        ]
+    }
+}
+
 def getDevelopmentVersion() {
     def output = new StringBuilder()
     def error = new StringBuilder()
@@ -79,9 +98,10 @@ dependencies {
 
     errorprone 'com.uber.nullaway:nullaway:0.12.6'
     errorprone 'com.google.errorprone:error_prone_core:2.37.0'
-}
 
-import net.ltgt.gradle.errorprone.CheckSeverity
+    // just tests
+    testCompileOnly 'org.jetbrains.kotlin:kotlin-stdlib-jdk8'
+}
 
 tasks.withType(JavaCompile) {
     options.release = 11
diff --git a/src/main/java/org/dataloader/DataLoader.java b/src/main/java/org/dataloader/DataLoader.java
index 7a50619..321b58c 100644
--- a/src/main/java/org/dataloader/DataLoader.java
+++ b/src/main/java/org/dataloader/DataLoader.java
@@ -68,7 +68,7 @@
  */
 @PublicApi
 @NullMarked
-public class DataLoader<K, V> {
+public class DataLoader<K, V extends @Nullable Object> {
 
     private final @Nullable String name;
     private final DataLoaderHelper<K, V> helper;
diff --git a/src/test/kotlin/org/dataloader/KotlinExamples.kt b/src/test/kotlin/org/dataloader/KotlinExamples.kt
new file mode 100644
index 0000000..da532e4
--- /dev/null
+++ b/src/test/kotlin/org/dataloader/KotlinExamples.kt
@@ -0,0 +1,39 @@
+package org.dataloader
+
+import org.junit.jupiter.api.Test
+import java.util.concurrent.CompletableFuture
+
+/**
+ * Some Kotlin code to prove that are JSpecify annotations work here
+ * as expected in Kotlin land.  We don't intend to ue Kotlin in our tests
+ * or to deliver Kotlin code in the java
+ */
+class KotlinExamples {
+
+    @Test
+    fun `basic kotlin test of non nullable value types`() {
+        val dataLoader: DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys -> CompletableFuture.completedFuture(keys.toList()) }
+
+        val cfA = dataLoader.load("A")
+        val cfB = dataLoader.load("B")
+
+        dataLoader.dispatch()
+
+        cfA.join().equals("A")
+        cfB.join().equals("B")
+    }
+
+    @Test
+    fun `basic kotlin test of nullable value types`() {
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newDataLoader { keys -> CompletableFuture.completedFuture(keys.toList()) }
+
+        val cfA = dataLoader.load("A")
+        val cfB = dataLoader.load("B")
+
+        dataLoader.dispatch()
+
+        cfA.join().equals("A")
+        cfB.join().equals("B")
+    }
+
+}
\ No newline at end of file

From 4a68a3930aae0f83ade11a2ccf70117a16f9ba38 Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 4 Jun 2025 15:23:44 +1000
Subject: [PATCH 03/15] Error Prone / NullAway support for JSpecify - added
 Kotlin - tweak

---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 2d22afd..551f4ec 100644
--- a/build.gradle
+++ b/build.gradle
@@ -16,7 +16,7 @@ plugins {
     id "me.champeau.jmh" version "0.7.3"
     id "net.ltgt.errorprone" version '4.2.0'
 
-    // Kotlin just for tests - not
+    // Kotlin just for tests - not production code
     id 'org.jetbrains.kotlin.jvm' version '2.1.21'
 }
 

From 7a6749e913e5ef7bcf281e15bc99d03a113b11b4 Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 4 Jun 2025 15:27:18 +1000
Subject: [PATCH 04/15] Error Prone / NullAway support for JSpecify - added
 Kotlin - tweak to assert

---
 src/test/kotlin/org/dataloader/KotlinExamples.kt | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/test/kotlin/org/dataloader/KotlinExamples.kt b/src/test/kotlin/org/dataloader/KotlinExamples.kt
index da532e4..c415b1a 100644
--- a/src/test/kotlin/org/dataloader/KotlinExamples.kt
+++ b/src/test/kotlin/org/dataloader/KotlinExamples.kt
@@ -2,6 +2,7 @@ package org.dataloader
 
 import org.junit.jupiter.api.Test
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.CompletableFuture.*
 
 /**
  * Some Kotlin code to prove that are JSpecify annotations work here
@@ -12,28 +13,28 @@ class KotlinExamples {
 
     @Test
     fun `basic kotlin test of non nullable value types`() {
-        val dataLoader: DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys -> CompletableFuture.completedFuture(keys.toList()) }
+        val dataLoader: DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys -> completedFuture(keys.toList()) }
 
         val cfA = dataLoader.load("A")
         val cfB = dataLoader.load("B")
 
         dataLoader.dispatch()
 
-        cfA.join().equals("A")
-        cfB.join().equals("B")
+        assert(cfA.join().equals("A"))
+        assert(cfA.join().equals("A"))
     }
 
     @Test
     fun `basic kotlin test of nullable value types`() {
-        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newDataLoader { keys -> CompletableFuture.completedFuture(keys.toList()) }
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newDataLoader { keys -> completedFuture(keys.toList()) }
 
         val cfA = dataLoader.load("A")
         val cfB = dataLoader.load("B")
 
         dataLoader.dispatch()
 
-        cfA.join().equals("A")
-        cfB.join().equals("B")
+        assert(cfA.join().equals("A"))
+        assert(cfA.join().equals("A"))
     }
 
 }
\ No newline at end of file

From afd5dc1c1b027242bc39549e7616860535dc9cad Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Thu, 5 Jun 2025 08:48:59 +1000
Subject: [PATCH 05/15] Error Prone / NullAway support for JSpecify - added
 Kotlin - reverted to old assertions

---
 src/main/java/org/dataloader/DataLoaderRegistry.java     | 3 +--
 src/test/java/org/dataloader/DataLoaderRegistryTest.java | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/main/java/org/dataloader/DataLoaderRegistry.java b/src/main/java/org/dataloader/DataLoaderRegistry.java
index 3b7409c..6bc79f6 100644
--- a/src/main/java/org/dataloader/DataLoaderRegistry.java
+++ b/src/main/java/org/dataloader/DataLoaderRegistry.java
@@ -143,8 +143,7 @@ private static DataLoaderOptions setInInstrumentation(DataLoaderOptions options,
      * @return this registry
      */
     public DataLoaderRegistry register(DataLoader<?, ?> dataLoader) {
-        String name = dataLoader.getName();
-        Objects.requireNonNull(name, "The DataLoader must have a non null name");
+        String name = Assertions.nonNull(dataLoader.getName(), () -> "The DataLoader must have a non null name");
         dataLoaders.put(name, nameAndInstrumentDL(name, instrumentation, dataLoader));
         return this;
     }
diff --git a/src/test/java/org/dataloader/DataLoaderRegistryTest.java b/src/test/java/org/dataloader/DataLoaderRegistryTest.java
index 4335375..89624d7 100644
--- a/src/test/java/org/dataloader/DataLoaderRegistryTest.java
+++ b/src/test/java/org/dataloader/DataLoaderRegistryTest.java
@@ -1,6 +1,5 @@
 package org.dataloader;
 
-import org.dataloader.impl.DataLoaderAssertionException;
 import org.dataloader.stats.SimpleStatisticsCollector;
 import org.dataloader.stats.Statistics;
 import org.junit.jupiter.api.Assertions;

From c6d87199de304a95306f897b96091eef787c4e3f Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 18 Jun 2025 12:10:51 +1000
Subject: [PATCH 06/15] Move to new Sonatype URLs

---
 .github/workflows/master.yml  | 2 ++
 .github/workflows/release.yml | 2 ++
 build.gradle                  | 9 ++++++---
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml
index 01b89bc..8643910 100644
--- a/.github/workflows/master.yml
+++ b/.github/workflows/master.yml
@@ -10,6 +10,8 @@ jobs:
     env:
       MAVEN_CENTRAL_USER: ${{ secrets.MAVEN_CENTRAL_USER }}
       MAVEN_CENTRAL_PASSWORD: ${{ secrets.MAVEN_CENTRAL_PASSWORD }}
+      MAVEN_CENTRAL_USER_NEW: ${{ secrets.MAVEN_CENTRAL_USER_NEW }}
+      MAVEN_CENTRAL_PASSWORD_NEW: ${{ secrets.MAVEN_CENTRAL_PASSWORD_NEW }}
       MAVEN_CENTRAL_PGP_KEY: ${{ secrets.MAVEN_CENTRAL_PGP_KEY }}
 
     steps:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index a574a68..fd614cc 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -14,6 +14,8 @@ jobs:
       MAVEN_CENTRAL_PGP_KEY: ${{ secrets.MAVEN_CENTRAL_PGP_KEY }}
       MAVEN_CENTRAL_USER: ${{ secrets.MAVEN_CENTRAL_USER }}
       MAVEN_CENTRAL_PASSWORD: ${{ secrets.MAVEN_CENTRAL_PASSWORD }}
+      MAVEN_CENTRAL_USER_NEW: ${{ secrets.MAVEN_CENTRAL_USER_NEW }}
+      MAVEN_CENTRAL_PASSWORD_NEW: ${{ secrets.MAVEN_CENTRAL_PASSWORD_NEW }}
       RELEASE_VERSION: ${{ github.event.inputs.version }}
 
     steps:
diff --git a/build.gradle b/build.gradle
index 551f4ec..80de227 100644
--- a/build.gradle
+++ b/build.gradle
@@ -224,9 +224,12 @@ publishing {
 nexusPublishing {
     repositories {
         sonatype {
-            username = System.env.MAVEN_CENTRAL_USER
-            password = System.env.MAVEN_CENTRAL_PASSWORD
-        }
+            username = System.env.MAVEN_CENTRAL_USER_NEW
+            password = System.env.MAVEN_CENTRAL_PASSWORD_NEW
+            // https://central.sonatype.org/publish/publish-portal-ossrh-staging-api/#configuration
+            nexusUrl.set(uri("https://ossrh-staging-api.central.sonatype.com/service/local/"))
+            // GraphQL Java does not publish snapshots, but adding this URL for completeness
+            snapshotRepositoryUrl.set(uri("https://central.sonatype.com/repository/maven-snapshots/"))        }
     }
 }
 

From d7f53ad2654a8b4fbc993b388bd28a820245e259 Mon Sep 17 00:00:00 2001
From: dondonz <13839920+dondonz@users.noreply.github.com>
Date: Sun, 22 Jun 2025 09:34:24 +1000
Subject: [PATCH 07/15] Update Gradle Github action

---
 .github/workflows/master.yml       | 2 +-
 .github/workflows/pull_request.yml | 2 +-
 .github/workflows/release.yml      | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml
index 8643910..ad3f9f2 100644
--- a/.github/workflows/master.yml
+++ b/.github/workflows/master.yml
@@ -16,7 +16,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: gradle/actions/wrapper-validation@v3
+      - uses: gradle/actions/wrapper-validation@v4
       - name: Set up JDK 11
         uses: actions/setup-java@v4
         with:
diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml
index f16bf96..573e169 100644
--- a/.github/workflows/pull_request.yml
+++ b/.github/workflows/pull_request.yml
@@ -14,7 +14,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: gradle/actions/wrapper-validation@v3
+      - uses: gradle/actions/wrapper-validation@v4
       - name: Set up JDK 11
         uses: actions/setup-java@v4
         with:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index fd614cc..64ef19c 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -20,7 +20,7 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
-      - uses: gradle/actions/wrapper-validation@v3
+      - uses: gradle/actions/wrapper-validation@v4
       - name: Set up JDK 11
         uses: actions/setup-java@v4
         with:

From dc948178890571faccc2a45468271ea667fad338 Mon Sep 17 00:00:00 2001
From: dondonz <13839920+dondonz@users.noreply.github.com>
Date: Sun, 22 Jun 2025 10:22:13 +1000
Subject: [PATCH 08/15] Add dependabot

---
 .github/dependabot.yml | 10 ++++++++++
 1 file changed, 10 insertions(+)
 create mode 100644 .github/dependabot.yml

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..10ef831
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,10 @@
+version: 2
+updates:
+  - package-ecosystem: "gradle"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"

From a1cd77c130d135dfe2b3252bf7bd9a0e994ea68f Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Sun, 22 Jun 2025 00:24:37 +0000
Subject: [PATCH 09/15] Bump com.gradle.develocity from 3.19 to 4.0.2

Bumps com.gradle.develocity from 3.19 to 4.0.2.

---
updated-dependencies:
- dependency-name: com.gradle.develocity
  dependency-version: 4.0.2
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 settings.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/settings.gradle b/settings.gradle
index 47404e7..d5a82f4 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -1,5 +1,5 @@
 plugins {
-    id 'com.gradle.develocity' version '3.19'
+    id 'com.gradle.develocity' version '4.0.2'
     id 'org.gradle.toolchains.foojay-resolver-convention' version '0.9.0'
 }
 

From 5764dffcea7cb89d07f4265573efd217f47806a6 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Sun, 22 Jun 2025 00:24:38 +0000
Subject: [PATCH 10/15] Bump com.github.ben-manes.versions from 0.51.0 to
 0.52.0

Bumps com.github.ben-manes.versions from 0.51.0 to 0.52.0.

---
updated-dependencies:
- dependency-name: com.github.ben-manes.versions
  dependency-version: 0.52.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 80de227..296c022 100644
--- a/build.gradle
+++ b/build.gradle
@@ -12,7 +12,7 @@ plugins {
     id 'groovy'
     id 'biz.aQute.bnd.builder' version '6.2.0'
     id 'io.github.gradle-nexus.publish-plugin' version '1.0.0'
-    id 'com.github.ben-manes.versions' version '0.51.0'
+    id 'com.github.ben-manes.versions' version '0.52.0'
     id "me.champeau.jmh" version "0.7.3"
     id "net.ltgt.errorprone" version '4.2.0'
 

From 262739b5b95e75372ba40aaeb7042620badd19c4 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Sun, 22 Jun 2025 00:24:41 +0000
Subject: [PATCH 11/15] Bump com.uber.nullaway:nullaway from 0.12.6 to 0.12.7

Bumps [com.uber.nullaway:nullaway](https://github.com/uber/NullAway) from 0.12.6 to 0.12.7.
- [Release notes](https://github.com/uber/NullAway/releases)
- [Changelog](https://github.com/uber/NullAway/blob/master/CHANGELOG.md)
- [Commits](https://github.com/uber/NullAway/compare/v0.12.6...v0.12.7)

---
updated-dependencies:
- dependency-name: com.uber.nullaway:nullaway
  dependency-version: 0.12.7
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 80de227..8a36148 100644
--- a/build.gradle
+++ b/build.gradle
@@ -96,7 +96,7 @@ dependencies {
     jmh 'org.openjdk.jmh:jmh-core:1.37'
     jmh 'org.openjdk.jmh:jmh-generator-annprocess:1.37'
 
-    errorprone 'com.uber.nullaway:nullaway:0.12.6'
+    errorprone 'com.uber.nullaway:nullaway:0.12.7'
     errorprone 'com.google.errorprone:error_prone_core:2.37.0'
 
     // just tests

From aa7349f7c4a44d16ae85a3c9ef10fb56c44b7c1d Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 25 Jun 2025 09:28:38 +1000
Subject: [PATCH 12/15] More explicit nullable annotations on BatchLoader
 interfaces

---
 src/main/java/org/dataloader/BatchLoader.java |  9 ++-
 .../dataloader/BatchLoaderWithContext.java    |  3 +-
 .../java/org/dataloader/DataLoaderHelper.java |  4 +-
 .../org/dataloader/MappedBatchLoader.java     |  3 +-
 .../MappedBatchLoaderWithContext.java         |  3 +-
 .../org/dataloader/MappedBatchPublisher.java  |  3 +-
 .../MappedBatchPublisherWithContext.java      |  3 +-
 .../kotlin/org/dataloader/KotlinExamples.kt   | 67 +++++++++++++++++--
 8 files changed, 77 insertions(+), 18 deletions(-)

diff --git a/src/main/java/org/dataloader/BatchLoader.java b/src/main/java/org/dataloader/BatchLoader.java
index 2b0c3c5..df11f89 100644
--- a/src/main/java/org/dataloader/BatchLoader.java
+++ b/src/main/java/org/dataloader/BatchLoader.java
@@ -17,8 +17,8 @@
 package org.dataloader;
 
 import org.dataloader.annotations.PublicSpi;
-import org.jspecify.annotations.NonNull;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 import java.util.List;
 import java.util.concurrent.CompletionStage;
@@ -40,7 +40,7 @@
  *      2, 9, 6, 1
  *  ]
  * </pre>
- *
+ * <p>
  * and loading from a back-end service returned this list of  values:
  *
  * <pre>
@@ -50,7 +50,7 @@
  *      { id: 2, name: 'San Francisco' },
  *  ]
  * </pre>
- *
+ * <p>
  * then the batch loader function contract has been broken.
  * <p>
  * The back-end service returned results in a different order than we requested, likely because it was more efficient for it to
@@ -77,7 +77,7 @@
 @FunctionalInterface
 @PublicSpi
 @NullMarked
-public interface BatchLoader<K, V> {
+public interface BatchLoader<K, V extends @Nullable Object> {
 
     /**
      * Called to batch load the provided keys and return a promise to a list of values.
@@ -85,7 +85,6 @@ public interface BatchLoader<K, V> {
      * If you need calling context then implement {@link org.dataloader.BatchLoaderWithContext}
      *
      * @param keys the collection of keys to load
-     *
      * @return a promise of the values for those keys
      */
     CompletionStage<List<V>> load(List<K> keys);
diff --git a/src/main/java/org/dataloader/BatchLoaderWithContext.java b/src/main/java/org/dataloader/BatchLoaderWithContext.java
index eba26e4..fb6ff71 100644
--- a/src/main/java/org/dataloader/BatchLoaderWithContext.java
+++ b/src/main/java/org/dataloader/BatchLoaderWithContext.java
@@ -2,6 +2,7 @@
 
 import org.dataloader.annotations.PublicSpi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 import java.util.List;
 import java.util.concurrent.CompletionStage;
@@ -16,7 +17,7 @@
  */
 @PublicSpi
 @NullMarked
-public interface BatchLoaderWithContext<K, V> {
+public interface BatchLoaderWithContext<K, V extends @Nullable Object> {
     /**
      * Called to batch load the provided keys and return a promise to a list of values.  This default
      * version can be given an environment object to that maybe be useful during the call.  A typical use case
diff --git a/src/main/java/org/dataloader/DataLoaderHelper.java b/src/main/java/org/dataloader/DataLoaderHelper.java
index 7858780..f4a8f10 100644
--- a/src/main/java/org/dataloader/DataLoaderHelper.java
+++ b/src/main/java/org/dataloader/DataLoaderHelper.java
@@ -594,11 +594,11 @@ private boolean isMapLoader() {
     }
 
     private boolean isPublisher() {
-        return batchLoadFunction instanceof BatchPublisher;
+        return batchLoadFunction instanceof BatchPublisher || batchLoadFunction instanceof BatchPublisherWithContext;
     }
 
     private boolean isMappedPublisher() {
-        return batchLoadFunction instanceof MappedBatchPublisher;
+        return batchLoadFunction instanceof MappedBatchPublisher || batchLoadFunction instanceof MappedBatchPublisherWithContext;
     }
 
     private DataLoaderInstrumentation instrumentation() {
diff --git a/src/main/java/org/dataloader/MappedBatchLoader.java b/src/main/java/org/dataloader/MappedBatchLoader.java
index 1ad4c79..179d6a2 100644
--- a/src/main/java/org/dataloader/MappedBatchLoader.java
+++ b/src/main/java/org/dataloader/MappedBatchLoader.java
@@ -18,6 +18,7 @@
 
 import org.dataloader.annotations.PublicSpi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 import java.util.Map;
 import java.util.Set;
@@ -59,7 +60,7 @@
  */
 @PublicSpi
 @NullMarked
-public interface MappedBatchLoader<K, V> {
+public interface MappedBatchLoader<K, V extends @Nullable Object> {
 
     /**
      * Called to batch load the provided keys and return a promise to a map of values.
diff --git a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java
index 9559260..9f342d4 100644
--- a/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java
+++ b/src/main/java/org/dataloader/MappedBatchLoaderWithContext.java
@@ -18,6 +18,7 @@
 
 import org.dataloader.annotations.PublicSpi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 
 import java.util.Map;
 import java.util.Set;
@@ -33,7 +34,7 @@
  */
 @PublicSpi
 @NullMarked
-public interface MappedBatchLoaderWithContext<K, V> {
+public interface MappedBatchLoaderWithContext<K, V extends @Nullable Object> {
     /**
      * Called to batch load the provided keys and return a promise to a map of values.
      *
diff --git a/src/main/java/org/dataloader/MappedBatchPublisher.java b/src/main/java/org/dataloader/MappedBatchPublisher.java
index 493401f..6637157 100644
--- a/src/main/java/org/dataloader/MappedBatchPublisher.java
+++ b/src/main/java/org/dataloader/MappedBatchPublisher.java
@@ -2,6 +2,7 @@
 
 import org.dataloader.annotations.PublicSpi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 import org.reactivestreams.Subscriber;
 
 import java.util.Map;
@@ -20,7 +21,7 @@
  */
 @PublicSpi
 @NullMarked
-public interface MappedBatchPublisher<K, V> {
+public interface MappedBatchPublisher<K, V extends @Nullable Object> {
     /**
      * Called to batch the provided keys into a stream of map entries of keys and values.
      * <p>
diff --git a/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java b/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java
index 7b862ca..dd8b5f9 100644
--- a/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java
+++ b/src/main/java/org/dataloader/MappedBatchPublisherWithContext.java
@@ -2,6 +2,7 @@
 
 import org.dataloader.annotations.PublicSpi;
 import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
 import org.reactivestreams.Subscriber;
 
 import java.util.List;
@@ -17,7 +18,7 @@
  */
 @PublicSpi
 @NullMarked
-public interface MappedBatchPublisherWithContext<K, V> {
+public interface MappedBatchPublisherWithContext<K, V extends @Nullable Object> {
 
     /**
      * Called to batch the provided keys into a stream of map entries of keys and values.
diff --git a/src/test/kotlin/org/dataloader/KotlinExamples.kt b/src/test/kotlin/org/dataloader/KotlinExamples.kt
index c415b1a..480a965 100644
--- a/src/test/kotlin/org/dataloader/KotlinExamples.kt
+++ b/src/test/kotlin/org/dataloader/KotlinExamples.kt
@@ -1,8 +1,8 @@
 package org.dataloader
 
 import org.junit.jupiter.api.Test
-import java.util.concurrent.CompletableFuture
-import java.util.concurrent.CompletableFuture.*
+import reactor.core.publisher.Flux
+import java.util.concurrent.CompletableFuture.completedFuture
 
 /**
  * Some Kotlin code to prove that are JSpecify annotations work here
@@ -13,7 +13,10 @@ class KotlinExamples {
 
     @Test
     fun `basic kotlin test of non nullable value types`() {
-        val dataLoader: DataLoader<String, String> = DataLoaderFactory.newDataLoader { keys -> completedFuture(keys.toList()) }
+        val batchLoadFunction = BatchLoader<String, String>
+        { keys -> completedFuture(keys.toList()) }
+        val dataLoader: DataLoader<String, String> =
+            DataLoaderFactory.newDataLoader(batchLoadFunction)
 
         val cfA = dataLoader.load("A")
         val cfB = dataLoader.load("B")
@@ -21,20 +24,72 @@ class KotlinExamples {
         dataLoader.dispatch()
 
         assert(cfA.join().equals("A"))
-        assert(cfA.join().equals("A"))
+        assert(cfB.join().equals("B"))
     }
 
     @Test
     fun `basic kotlin test of nullable value types`() {
-        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newDataLoader { keys -> completedFuture(keys.toList()) }
+        val batchLoadFunction: BatchLoader<String, String?> = BatchLoader { keys -> completedFuture(keys.toList()) }
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newDataLoader(batchLoadFunction)
+
+        standardNullableAsserts(dataLoader)
+    }
+
+    @Test
+    fun `basic kotlin test of nullable value types in mapped batch loader`() {
+        val batchLoadFunction = MappedBatchLoader<String, String?>
+        { keys -> completedFuture(keys.associateBy({ it })) }
+
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newMappedDataLoader(batchLoadFunction)
+
+        standardNullableAsserts(dataLoader)
+    }
+
+    @Test
+    fun `basic kotlin test of nullable value types in mapped batch loader with context`() {
+        val batchLoadFunction = MappedBatchLoaderWithContext<String, String?>
+        { keys, env -> completedFuture(keys.associateBy({ it })) }
+
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newMappedDataLoader(batchLoadFunction)
+
+        standardNullableAsserts(dataLoader)
+    }
 
+    @Test
+    fun `basic kotlin test of nullable value types in mapped batch publisher`() {
+        val batchLoadFunction = MappedBatchPublisher<String, String?>
+        { keys, subscriber ->
+            val map: Map<String, String?> = keys.associateBy({ it })
+            Flux.fromIterable(map.entries).subscribe(subscriber);
+        }
+
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newMappedPublisherDataLoader(batchLoadFunction)
+
+        standardNullableAsserts(dataLoader)
+    }
+
+    @Test
+    fun `basic kotlin test of nullable value types in mapped batch publisher with context`() {
+        val batchLoadFunction = MappedBatchPublisherWithContext<String, String?>
+        { keys, subscriber, env ->
+            val map: Map<String, String?> = keys.associateBy({ it })
+            Flux.fromIterable(map.entries).subscribe(subscriber);
+        }
+
+        val dataLoader: DataLoader<String, String?> = DataLoaderFactory.newMappedPublisherDataLoader(batchLoadFunction)
+
+        standardNullableAsserts(dataLoader)
+    }
+
+    private fun standardNullableAsserts(dataLoader: DataLoader<String, String?>) {
         val cfA = dataLoader.load("A")
         val cfB = dataLoader.load("B")
 
         dataLoader.dispatch()
 
         assert(cfA.join().equals("A"))
-        assert(cfA.join().equals("A"))
+        assert(cfB.join().equals("B"))
     }
 
+
 }
\ No newline at end of file

From e10797eff8ef6e4fcf75e3fe60e8c4a5104d157a Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 30 Jun 2025 13:33:29 +0000
Subject: [PATCH 13/15] Bump org.jetbrains.kotlin.jvm from 2.1.21 to 2.2.0

Bumps [org.jetbrains.kotlin.jvm](https://github.com/JetBrains/kotlin) from 2.1.21 to 2.2.0.
- [Release notes](https://github.com/JetBrains/kotlin/releases)
- [Changelog](https://github.com/JetBrains/kotlin/blob/master/ChangeLog.md)
- [Commits](https://github.com/JetBrains/kotlin/compare/v2.1.21...v2.2.0)

---
updated-dependencies:
- dependency-name: org.jetbrains.kotlin.jvm
  dependency-version: 2.2.0
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 8a36148..5ed119a 100644
--- a/build.gradle
+++ b/build.gradle
@@ -17,7 +17,7 @@ plugins {
     id "net.ltgt.errorprone" version '4.2.0'
 
     // Kotlin just for tests - not production code
-    id 'org.jetbrains.kotlin.jvm' version '2.1.21'
+    id 'org.jetbrains.kotlin.jvm' version '2.2.0'
 }
 
 java {

From e6795fe988487b237cfc924937d4e526652310f4 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 1 Jul 2025 22:48:18 +0000
Subject: [PATCH 14/15] Bump io.github.gradle-nexus.publish-plugin from 1.0.0
 to 2.0.0

Bumps io.github.gradle-nexus.publish-plugin from 1.0.0 to 2.0.0.

---
updated-dependencies:
- dependency-name: io.github.gradle-nexus.publish-plugin
  dependency-version: 2.0.0
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 build.gradle | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build.gradle b/build.gradle
index 065f745..3edf4f8 100644
--- a/build.gradle
+++ b/build.gradle
@@ -11,7 +11,7 @@ plugins {
     id 'signing'
     id 'groovy'
     id 'biz.aQute.bnd.builder' version '6.2.0'
-    id 'io.github.gradle-nexus.publish-plugin' version '1.0.0'
+    id 'io.github.gradle-nexus.publish-plugin' version '2.0.0'
     id 'com.github.ben-manes.versions' version '0.52.0'
     id "me.champeau.jmh" version "0.7.3"
     id "net.ltgt.errorprone" version '4.2.0'

From 1d34141834b0ec89b663bbb6d76df36c7c227805 Mon Sep 17 00:00:00 2001
From: bbaker <bbaker@atlassian.com>
Date: Wed, 2 Jul 2025 16:29:20 +1000
Subject: [PATCH 15/15] #208 to ensure tests are in place

---
 .../java/org/dataloader/DataLoaderTest.java   | 26 +++++++++++++++++++
 .../DelegatingDataLoaderFactory.java          |  7 +++++
 .../parameterized/ListDataLoaderFactory.java  | 11 ++++++++
 .../MappedDataLoaderFactory.java              | 14 +++++++++-
 .../MappedPublisherDataLoaderFactory.java     | 18 ++++++++++---
 .../PublisherDataLoaderFactory.java           | 12 ++++++++-
 .../parameterized/TestDataLoaderFactory.java  |  7 +++++
 7 files changed, 89 insertions(+), 6 deletions(-)

diff --git a/src/test/java/org/dataloader/DataLoaderTest.java b/src/test/java/org/dataloader/DataLoaderTest.java
index 224b54d..37ae030 100644
--- a/src/test/java/org/dataloader/DataLoaderTest.java
+++ b/src/test/java/org/dataloader/DataLoaderTest.java
@@ -45,6 +45,7 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -1230,6 +1231,31 @@ public void when_values_size_are_more_then_key_size(TestDataLoaderFactory factor
         }
     }
 
+    @ParameterizedTest
+    @MethodSource("org.dataloader.fixtures.parameterized.TestDataLoaderFactories#get")
+    public void should_Support_loading_values_with_context(TestDataLoaderFactory factory) {
+        AtomicReference<BatchLoaderEnvironment> environmentREF = new AtomicReference<>();
+        DataLoader<Integer, Integer> identityLoader = factory.idLoaderWithContext(new DataLoaderOptions(), new ArrayList<>(), environmentREF);
+
+        identityLoader.load(1, "ctx1");
+        identityLoader.load(2, "ctx2");
+        identityLoader.loadMany(List.of(3, 4), List.of("ctx3", "ctx4"));
+
+        CompletableFuture<List<Integer>> cf = identityLoader.dispatch();
+        await().atMost(Duration.FIVE_SECONDS).until(cf::isDone);
+
+        assertThat(cf.toCompletableFuture().join(), equalTo(asList(1, 2, 3, 4)));
+
+        Map<Object, Object> keyContexts = environmentREF.get().getKeyContexts();
+        assertThat(keyContexts, equalTo(Map.of(
+                1, "ctx1",
+                2, "ctx2",
+                3, "ctx3",
+                4, "ctx4"
+        )));
+    }
+
+
     @Test
     public void can_call_a_loader_from_a_loader() throws Exception {
         List<Collection<String>> deepLoadCalls = new ArrayList<>();
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java
index 0cbd3f3..8d1f815 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/DelegatingDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.DelegatingDataLoader;
@@ -8,6 +9,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class DelegatingDataLoaderFactory implements TestDataLoaderFactory {
     // its delegates all the way down to the turtles
@@ -38,6 +40,11 @@ public <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<
         return mkDelegateDataLoader(delegateFactory.idLoader(options, loadCalls));
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return mkDelegateDataLoader(delegateFactory.idLoaderWithContext(options, loadCalls, environmentREF));
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return mkDelegateDataLoader(delegateFactory.idLoaderDelayed(options, loadCalls, delay));
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java
index 0644d3c..8ec69d7 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/ListDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.fixtures.TestKit;
@@ -9,6 +10,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
@@ -23,6 +25,15 @@ public <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newDataLoader((keys, env) -> {
+            environmentREF.set(env);
+            loadCalls.add(new ArrayList<>(keys));
+            return completedFuture(keys);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return newDataLoader(keys -> CompletableFuture.supplyAsync(() -> {
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/MappedDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/MappedDataLoaderFactory.java
index e7c47ec..f1c548e 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/MappedDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/MappedDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.fixtures.TestKit;
@@ -11,10 +12,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 
 import static java.util.concurrent.CompletableFuture.completedFuture;
-import static org.dataloader.DataLoaderFactory.newDataLoader;
 import static org.dataloader.DataLoaderFactory.newMappedDataLoader;
 import static org.dataloader.fixtures.TestKit.futureError;
 
@@ -31,6 +32,17 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newMappedDataLoader((keys, environment) -> {
+            environmentREF.set(environment);
+            loadCalls.add(new ArrayList<>(keys));
+            Map<K, K> map = new HashMap<>();
+            keys.forEach(k -> map.put(k, k));
+            return completedFuture(map);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(
             DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/MappedPublisherDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/MappedPublisherDataLoaderFactory.java
index fa920cf..3a0f54e 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/MappedPublisherDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/MappedPublisherDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.Try;
@@ -12,16 +13,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.CompletableFuture;
-import java.util.stream.Collectors;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toList;
-import static java.util.stream.Collectors.toSet;
 import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoader;
 import static org.dataloader.DataLoaderFactory.newMappedPublisherDataLoaderWithTry;
-import static org.dataloader.DataLoaderFactory.newPublisherDataLoader;
 
 public class MappedPublisherDataLoaderFactory implements TestDataLoaderFactory, TestReactiveDataLoaderFactory {
 
@@ -36,6 +34,18 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newMappedPublisherDataLoader((keys, subscriber, environment) -> {
+            environmentREF.set(environment);
+
+            loadCalls.add(new ArrayList<>(keys));
+            Map<K, K> map = new HashMap<>();
+            keys.forEach(k -> map.put(k, k));
+            Flux.fromIterable(map.entrySet()).subscribe(subscriber);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(
             DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/PublisherDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/PublisherDataLoaderFactory.java
index 2049719..c8e8b67 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/PublisherDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/PublisherDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 import org.dataloader.Try;
@@ -11,9 +12,9 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
-import static org.dataloader.DataLoaderFactory.newDataLoader;
 import static org.dataloader.DataLoaderFactory.newPublisherDataLoader;
 import static org.dataloader.DataLoaderFactory.newPublisherDataLoaderWithTry;
 
@@ -28,6 +29,15 @@ public <K> DataLoader<K, K> idLoader(
         }, options);
     }
 
+    @Override
+    public <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF) {
+        return newPublisherDataLoader((keys, subscriber, environment) -> {
+            environmentREF.set(environment);
+            loadCalls.add(new ArrayList<>(keys));
+            Flux.fromIterable(keys).subscribe(subscriber);
+        }, options);
+    }
+
     @Override
     public <K> DataLoader<K, K> idLoaderDelayed(DataLoaderOptions options, List<Collection<K>> loadCalls, Duration delay) {
         return newPublisherDataLoader((keys, subscriber) -> {
diff --git a/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java b/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java
index 789b136..3c584fd 100644
--- a/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java
+++ b/src/test/java/org/dataloader/fixtures/parameterized/TestDataLoaderFactory.java
@@ -1,5 +1,6 @@
 package org.dataloader.fixtures.parameterized;
 
+import org.dataloader.BatchLoaderEnvironment;
 import org.dataloader.DataLoader;
 import org.dataloader.DataLoaderOptions;
 
@@ -7,6 +8,7 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 
 public interface TestDataLoaderFactory {
     <K> DataLoader<K, K> idLoader(DataLoaderOptions options, List<Collection<K>> loadCalls);
@@ -23,6 +25,11 @@ public interface TestDataLoaderFactory {
 
     DataLoader<String, String> idLoaderReturnsTooMany(int howManyMore, DataLoaderOptions options, ArrayList<Object> loadCalls);
 
+    // similar to above but batch loaders with context
+
+    <K> DataLoader<K, K> idLoaderWithContext(DataLoaderOptions options, List<Collection<K>> loadCalls, AtomicReference<BatchLoaderEnvironment> environmentREF);
+
+
     // Convenience methods
 
     default <K> DataLoader<K, K> idLoader(DataLoaderOptions options) {

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://rainy.clevelandohioweatherforecast.com/php-proxy/index.php?q=https://github.com/graphql-java/java-dataloader/compare/v5.0.0...v5.0.1.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/index.php?u=https://github.com/graphql-java/java-dataloader/compare/v5.0.0...v5.0.1.patch" target="_blank">pFad Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v3index.php?u=https://github.com/graphql-java/java-dataloader/compare/v5.0.0...v5.0.1.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v4index.php?u=https://github.com/graphql-java/java-dataloader/compare/v5.0.0...v5.0.1.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>