fix: handle use case where data fetched is null to avoid exceptions

Signed-off-by: TessaIO <ahmedgrati1999@gmail.com>
apache · Mar 19, 2024 · a646858 · a646858
1 parent 8e417d2
commit a646858
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 24 deletions.
diff --git a/...ore/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java b/...ore/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java
@@ -27,11 +27,7 @@
 import org.apache.druid.server.lookup.cache.loading.LoadingCache;
 
 import javax.annotation.Nullable;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.HashSet;
+import java.util.*;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -74,16 +70,20 @@ public String apply(@Nullable final String key)
       // otherwise null will be replaced with empty string in nullToEmptyIfNeeded above.
       return null;
     }
-
-    final String presentVal;
-    try {
-      presentVal = loadingCache.get(keyEquivalent, new ApplyCallable(keyEquivalent));
+    String presentVal = 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);
+    LOGGER.info("VAL is [%s] AND KEY IS [%s]", val, keyEquivalent);
+    if (val == null) {
       return null;
     }
+    Map<String, String> map = new HashMap<String, String>();
+    map.put(keyEquivalent, val);
+    loadingCache.putAll(map);
+    return NullHandling.emptyToNullIfNeeded(val);
   }
 
   @Override

diff --git a/...lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java b/...lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java
@@ -48,7 +48,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(""));
@@ -74,7 +74,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);
@@ -101,17 +101,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
   {