From 94b1c904e88730b88376c1ec8da77d41220c6326 Mon Sep 17 00:00:00 2001
From: David Leifker <david.leifker@acryl.io>
Date: Tue, 3 Sep 2024 10:45:52 -0500
Subject: [PATCH] feat(graphql): Lazy dataLoaders

---
 .../datahub/graphql/GraphQLEngine.java        | 13 +----
 .../graphql/LazyDataLoaderRegistry.java       | 48 +++++++++----------
 2 files changed, 24 insertions(+), 37 deletions(-)

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java
index dd8eabd3ce06fd..97e282c106a9a2 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GraphQLEngine.java
@@ -28,7 +28,6 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.dataloader.DataLoader;
-import org.dataloader.DataLoaderRegistry;
 
 /**
  * Simple wrapper around a {@link GraphQL} instance providing APIs for building an engine and
@@ -100,7 +99,7 @@ public ExecutionResult execute(
     /*
      * Init DataLoaderRegistry - should be created for each request.
      */
-    DataLoaderRegistry register = createDataLoaderRegistry(_dataLoaderSuppliers, context);
+    LazyDataLoaderRegistry register = new LazyDataLoaderRegistry(context, _dataLoaderSuppliers);
 
     /*
      * Construct execution input
@@ -218,14 +217,4 @@ public GraphQLEngine build() {
           graphQLQueryIntrospectionEnabled);
     }
   }
-
-  private DataLoaderRegistry createDataLoaderRegistry(
-      final Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers,
-      final QueryContext context) {
-    final DataLoaderRegistry registry = new DataLoaderRegistry();
-    for (String key : dataLoaderSuppliers.keySet()) {
-      registry.register(key, dataLoaderSuppliers.get(key).apply(context));
-    }
-    return registry;
-  }
 }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java
index bd668c5810bdb3..2a0b70aa07cc62 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/LazyDataLoaderRegistry.java
@@ -1,34 +1,32 @@
 package com.linkedin.datahub.graphql;
 
-import org.dataloader.DataLoader;
-import org.dataloader.DataLoaderRegistry;
-
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Function;
-import java.util.function.Supplier;
-
-public class LazyDataLoaderManager {
-    private final Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers = new ConcurrentHashMap<>();
-    private final DataLoaderRegistry registry;
+import org.dataloader.DataLoader;
+import org.dataloader.DataLoaderRegistry;
 
-    public LazyDataLoaderManager(Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers) {
-        this.dataLoaderSuppliers.putAll(dataLoaderSuppliers);
-        this.registry = registry;
-    }
+public class LazyDataLoaderRegistry extends DataLoaderRegistry {
+  private final QueryContext queryContext;
+  private final Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers;
+  private final DataLoaderRegistry registry;
 
-    public <K, V> DataLoader<K, V> getDataLoader(String key) {
-        if (!registry.getKeys().contains(key)) {
-            Function<QueryContext, DataLoader<?, ?>> supplier = dataLoaderSuppliers.get(key);
-            if (supplier == null) {
-                throw new IllegalArgumentException("No DataLoader registered for key: " + key);
-            }
-            registry.register(key, supplier.get());
-        }
-        return registry.getDataLoader(key);
-    }
+  public LazyDataLoaderRegistry(
+      QueryContext queryContext,
+      Map<String, Function<QueryContext, DataLoader<?, ?>>> dataLoaderSuppliers) {
+    this.queryContext = queryContext;
+    this.dataLoaderSuppliers = dataLoaderSuppliers;
+    this.registry = new DataLoaderRegistry();
+  }
 
-    public void clearAll() {
-        registry.getKeys().forEach(registry::unregister);
+  @Override
+  public <K, V> DataLoader<K, V> getDataLoader(String key) {
+    if (!registry.getKeys().contains(key)) {
+      Function<QueryContext, DataLoader<?, ?>> supplier = dataLoaderSuppliers.get(key);
+      if (supplier == null) {
+        throw new IllegalArgumentException("No DataLoader registered for key: " + key);
+      }
+      registry.register(key, supplier.apply(queryContext));
     }
+    return registry.getDataLoader(key);
+  }
 }