Use groupedKeys to do N+1 query

src/main/java/com/longfor/bff_netflix/customcontext/ExtendElementQueryArgument.java

@@ -0,0 +1,8 @@
+package com.longfor.bff_netflix.customcontext;
+
+import lombok.Data;
+
+@Data
+public class ExtendElementQueryArgument {
+    private String type;
+}

src/main/java/com/longfor/bff_netflix/customcontext/ExtendElementQueryKey.java

@@ -0,0 +1,9 @@
+package com.longfor.bff_netflix.customcontext;
+
+import lombok.Data;
+
+@Data
+public class ExtendElementQueryKey {
+    private Integer id;
+    private ExtendElementQueryArgument arguments;
+}

src/main/java/com/longfor/bff_netflix/datafetchers/ExtendElementDataFetcher.java

@@ -1,12 +1,12 @@
 package com.longfor.bff_netflix.datafetchers;
 
 import com.longfor.DgsConstants;
+import com.longfor.bff_netflix.customcontext.ExtendElementQueryArgument;
+import com.longfor.bff_netflix.customcontext.ExtendElementQueryKey;
 import com.longfor.bff_netflix.customcontext.ShopCustomContext;
-import com.longfor.bff_netflix.dataloaders.BaseShopInfoDataLoader;
+import com.longfor.bff_netflix.dataloaders.ExtendElementDataLoaderWithCustomContext;
-import com.longfor.bff_netflix.dataloaders.ExtendElementDataLoader;
+import com.longfor.bff_netflix.dataloaders.ExtendElementDataLoaderWithMappedKeys;
 import com.longfor.bff_netflix.services.ExtendElementService;
-import com.longfor.bff_netflix.services.ShopService;
-import com.longfor.types.BaseShopInfo;
 import com.longfor.types.ExtendElement;
 import com.longfor.types.Shop;
 import com.netflix.graphql.dgs.DgsComponent;
@@ -31,15 +31,29 @@ public class ExtendElementDataFetcher {
      * This datafetcher resolves the shows field on Query.
      * It uses an @InputArgument to get the titleFilter from the Query if one is defined.
      */
-    @DgsData(parentType = DgsConstants.SHOP.TYPE_NAME, field = DgsConstants.SHOP.ExpandData)
+    @DgsData(parentType = DgsConstants.SHOP.TYPE_NAME, field = DgsConstants.SHOP.ExtendElements)
     public CompletableFuture<List<ExtendElement>> extendElements(@InputArgument("typeNameFilter") String typeNameFilter , DgsDataFetchingEnvironment dfe) {
+        // method 1: Use custom context to pass the arguments to data loader
+//        Shop shop = dfe.getSource();
+//        DataLoader<Integer, List<ExtendElement>> dataLoader = dfe.getDataLoader(ExtendElementDataLoaderWithCustomContext.class);
+//        ShopCustomContext customContext =  DgsContext.getCustomContext(dfe);
+//        customContext.setExtendElementTypeName(typeNameFilter);
+//        return dataLoader.load(shop.getShopId());
 
+        // method 2: Use the mapped keys to pass the arguments to data loader
+        // 2.1 As passed ExtendElementQueryKey as key and dataloader accepts the Set<ExtendElementQueryKey, R> as parameters,
+        // Let's say we have keyA, keyB, keyC, and keyA holds the same value with keyB, then the downstream DataLoader will only have 2
+        // EntrySet, the query will be only 2
         Shop shop = dfe.getSource();
-        DataLoader<Integer, List<ExtendElement>> dataLoader = dfe.getDataLoader(ExtendElementDataLoader.class);
+        DataLoader<ExtendElementQueryKey, List<ExtendElement>> dataLoader = dfe.getDataLoader(ExtendElementDataLoaderWithMappedKeys.class);
-        ShopCustomContext customContext =  DgsContext.getCustomContext(dfe);
+        ExtendElementQueryArgument arg = new ExtendElementQueryArgument();
-        customContext.setExtendElementTypeName(typeNameFilter);
+        arg.setType(dfe.getExecutionStepInfo().getArgument("typeNameFilter"));
+        ExtendElementQueryKey queryKey = new ExtendElementQueryKey();
+        queryKey.setId(shop.getShopId());
+        queryKey.setArguments(arg);
+
+        return dataLoader.load(queryKey);
 
-        return dataLoader.load(shop.getShopId());
     }
 
 

src/main/java/com/longfor/bff_netflix/dataloaders/ExtendElementDataLoaderWithCustomContext.java

@@ -4,15 +4,10 @@ import com.longfor.DgsConstants;
 import com.longfor.bff_netflix.customcontext.ShopCustomContext;
 import com.longfor.bff_netflix.services.ExtendElementService;
 import com.longfor.types.ExtendElement;
-import com.netflix.graphql.dgs.DgsDataFetchingEnvironment;
 import com.netflix.graphql.dgs.DgsDataLoader;
 import com.netflix.graphql.dgs.context.DgsContext;
-import graphql.execution.DataFetcherExceptionHandler;
-import graphql.schema.DataFetchingEnvironment;
 import org.dataloader.BatchLoaderEnvironment;
-import org.dataloader.MappedBatchLoader;
 import org.dataloader.MappedBatchLoaderWithContext;
-import reactor.util.function.Tuple2;
 
 import java.util.List;
 import java.util.Map;
@@ -21,12 +16,12 @@ import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import java.util.stream.Collectors;
 
-@DgsDataLoader(name = DgsConstants.SHOP.ExpandData)
+@DgsDataLoader(name = DgsConstants.SHOP.ExtendElements)
-public class ExtendElementDataLoader implements MappedBatchLoaderWithContext<Integer, List<ExtendElement>> {
+public class ExtendElementDataLoaderWithCustomContext implements MappedBatchLoaderWithContext<Integer, List<ExtendElement>> {
     private final ExtendElementService extendElementService;
 
 
-    public ExtendElementDataLoader(ExtendElementService extendElementService) {
+    public ExtendElementDataLoaderWithCustomContext(ExtendElementService extendElementService) {
         this.extendElementService = extendElementService;
     }
 

src/main/java/com/longfor/bff_netflix/dataloaders/ExtendElementDataLoaderWithMappedKeys.java

@@ -0,0 +1,60 @@
+package com.longfor.bff_netflix.dataloaders;
+
+import com.longfor.DgsConstants;
+import com.longfor.bff_netflix.customcontext.ExtendElementQueryArgument;
+import com.longfor.bff_netflix.customcontext.ExtendElementQueryKey;
+import com.longfor.bff_netflix.customcontext.ShopCustomContext;
+import com.longfor.bff_netflix.services.ExtendElementService;
+import com.longfor.types.ExtendElement;
+import com.netflix.graphql.dgs.DgsDataLoader;
+import com.netflix.graphql.dgs.context.DgsContext;
+import org.dataloader.BatchLoaderEnvironment;
+import org.dataloader.MappedBatchLoaderWithContext;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
+import java.util.stream.Collectors;
+
+/**
+ * Showcase the query:
+ * query Shops {
+ *     shops {
+ *         shopId
+ *         baseShopInfo {
+ *             shopId
+ *             shopName
+ *             shopStatus
+ *         }
+ *         expandData:  extendElements   (typeNameFilter: "TypeA"){
+ *             attId,code,type,value
+ *         }
+ *         expandData2: extendElements(typeNameFilter: "TypeB"){
+ *             attId,code,type,value
+ *         }
+ *     }
+ * }
+ */
+@DgsDataLoader(name = DgsConstants.SHOP.ExtendElements)
+public class ExtendElementDataLoaderWithMappedKeys implements MappedBatchLoaderWithContext<ExtendElementQueryKey, List<ExtendElement>> {
+    private final ExtendElementService extendElementService;
+
+
+    public ExtendElementDataLoaderWithMappedKeys(ExtendElementService extendElementService) {
+        this.extendElementService = extendElementService;
+    }
+
+    @Override
+    public CompletionStage<Map<ExtendElementQueryKey, List<ExtendElement>>> load(Set<ExtendElementQueryKey> set, BatchLoaderEnvironment environment) {
+
+
+        return CompletableFuture.supplyAsync(() -> {
+            // groupBy arguments, so the query will hold the same argument with different shopId
+            Map<ExtendElementQueryArgument, List<ExtendElementQueryKey>> queryMap = set.stream().collect(Collectors.groupingBy(ExtendElementQueryKey::getArguments));
+           return  null;
+
+        });
+    }
+}

src/main/resources/schema/schema.graphqls

@@ -12,7 +12,7 @@ type Contract {
 type Shop{
     shopId: Int
     baseShopInfo:BaseShopInfo
-    expandData(typeNameFilter: String): [ExtendElement!]!
+    extendElements(typeNameFilter: String): [ExtendElement]
 }
 
 type BaseShopInfo{