From d9e4cd96c97874026c5a19f3cc3df6ac5b7d8206 Mon Sep 17 00:00:00 2001
From: TessaIO <ahmedgrati1999@gmail.com>
Date: Tue, 1 Oct 2024 16:46:54 +0200
Subject: [PATCH] fix loading lookup extension

Signed-off-by: TessaIO <ahmedgrati1999@gmail.com>
---
 .../druid/server/lookup/LoadingLookup.java    | 26 ++++++++++++++-----
 .../server/lookup/LoadingLookupTest.java      | 26 +++++++------------
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java b/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java
index 508b07b9330e2..9af52bcf3ea2b 100644
--- a/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java
+++ b/extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java
@@ -28,6 +28,7 @@
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Callable;
@@ -74,14 +75,22 @@ public String apply(@Nullable final String key)
     }
 
     final String presentVal;
-    try {
-      presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent));
+    presentVal = this.loadingCache.getIfPresent(keyEquivalent);
+    if (presentVal != null) {
       return NullHandling.emptyToNullIfNeeded(presentVal);
     }
-    catch (ExecutionException e) {
-      LOGGER.debug("value not found for key [%s]", key);
+
+    String val = this.dataFetcher.fetch(keyEquivalent);
+    if (val == null) {
       return null;
     }
+
+    Map<String, String> map = new HashMap<>();
+    map.put(keyEquivalent, val);
+
+    this.loadingCache.putAll(map);
+
+    return NullHandling.emptyToNullIfNeeded(val);
   }
 
   @Override
@@ -108,13 +117,18 @@ public List<String> unapply(@Nullable final String value)
   @Override
   public boolean supportsAsMap()
   {
-    return false;
+    return true;
   }
 
   @Override
   public Map<String, String> asMap()
   {
-    throw new UnsupportedOperationException("Cannot get map view");
+    Iterable<Map.Entry<String, String>> data = this.dataFetcher.fetchAll();
+    Map<String, String> map = new HashMap<>();
+    if (data != null) {
+      data.forEach(each -> map.put(each.getKey(), each.getValue()));
+    }
+    return map;
   }
 
   @Override
diff --git a/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java b/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java
index 43588bf8c4765..e96df27f3947b 100644
--- a/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java
+++ b/extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java
@@ -31,8 +31,10 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
+import java.util.AbstractMap;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 
@@ -49,7 +51,7 @@ public class LoadingLookupTest extends InitializedNullHandlingTest
   @Test
   public void testApplyEmptyOrNull() throws ExecutionException
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq(""), EasyMock.anyObject(Callable.class)))
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("")))
             .andReturn("empty").atLeastOnce();
     EasyMock.replay(lookupCache);
     Assert.assertEquals("empty", loadingLookup.apply(""));
@@ -75,7 +77,7 @@ public void testUnapplyNull()
   @Test
   public void testApply() throws ExecutionException
   {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class))).andReturn("value").once();
+    EasyMock.expect(lookupCache.getIfPresent(EasyMock.eq("key"))).andReturn("value").once();
     EasyMock.replay(lookupCache);
     Assert.assertEquals(ImmutableMap.of("key", "value"), loadingLookup.applyAll(ImmutableSet.of("key")));
     EasyMock.verify(lookupCache);
@@ -105,17 +107,6 @@ public void testClose()
     EasyMock.verify(lookupCache, reverseLookupCache);
   }
 
-  @Test
-  public void testApplyWithExecutionError() throws ExecutionException
-  {
-    EasyMock.expect(lookupCache.get(EasyMock.eq("key"), EasyMock.anyObject(Callable.class)))
-            .andThrow(new ExecutionException(null))
-            .once();
-    EasyMock.replay(lookupCache);
-    Assert.assertNull(loadingLookup.apply("key"));
-    EasyMock.verify(lookupCache);
-  }
-
   @Test
   public void testUnApplyWithExecutionError() throws ExecutionException
   {
@@ -136,13 +127,16 @@ public void testGetCacheKey()
   @Test
   public void testSupportsAsMap()
   {
-    Assert.assertFalse(loadingLookup.supportsAsMap());
+    Assert.assertTrue(loadingLookup.supportsAsMap());
   }
 
   @Test
   public void testAsMap()
   {
-    expectedException.expect(UnsupportedOperationException.class);
-    loadingLookup.asMap();
+    Map.Entry<String, String> fetchedData = new AbstractMap.SimpleEntry<>("dummy", "test");
+    EasyMock.expect(dataFetcher.fetchAll()).andReturn(Collections.singletonList(fetchedData));
+    EasyMock.replay(dataFetcher);
+    Assert.assertEquals(loadingLookup.asMap().size(), 1);
+    EasyMock.verify(dataFetcher);
   }
 }