Merge pull request #207 from graphql-java/jspecify-nullable-on-batch-loader-methods

More explicit nullable annotations on BatchLoader interfaces

Author: bbakerman

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,15 +77,14 @@
 @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.
      * <p>
      * 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);

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

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() {

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.

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.
      *

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>

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.

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,28 +13,83 @@ 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")
 
         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"))
     }
 
+
 }