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

Signed-off-by: TessaIO <ahmedgrati1999@gmail.com>

Author: TessaIO

extensions-core/lookups-cached-single/src/main/java/org/apache/druid/server/lookup/LoadingLookup.java

@@ -28,10 +28,11 @@
 
 import javax.annotation.Nullable;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.HashSet;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -74,16 +75,19 @@ 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);
+    if (val == null) {
       return null;
     }
+    Map<String, String> map = new HashMap<>();
+    map.put(keyEquivalent, val);
+    loadingCache.putAll(map);
+    return NullHandling.emptyToNullIfNeeded(val);
   }
 
   @Override
@@ -131,9 +135,7 @@ public Set<String> keySet()
     Iterable<Map.Entry<String, String>> data = this.dataFetcher.fetchAll();
     Set<String> set = new HashSet<>();
     if (data != null) {
-      data.forEach(each -> {
-        set.add(each.getKey());
-      });
+      data.forEach(each -> set.add(each.getKey()));
     }
     return set;
   }

extensions-core/lookups-cached-single/src/test/java/org/apache/druid/server/lookup/LoadingLookupTest.java

@@ -19,7 +19,6 @@
 
 package org.apache.druid.server.lookup;
 
-import com.amazonaws.transform.MapEntry;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import org.apache.druid.common.config.NullHandling;
@@ -31,7 +30,11 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
-import java.util.*;
+import java.util.AbstractMap;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.Map;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 
@@ -48,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(""));
@@ -74,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);
@@ -101,17 +104,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
   {
@@ -161,9 +153,13 @@ public void testFetchAll()
     EasyMock.verify(dataFetcher);
   }
 
-  public int getIteratorSize(Iterator<Map.Entry<String, String>> it) {
-    int i = 0;
-    for ( ; it.hasNext() ; ++i ) it.next();
-    return i;
+  public int getIteratorSize(Iterator<Map.Entry<String, String>> it)
+  {
+    int sum = 0;
+    while (it.hasNext()) {
+      sum++;
+      it.next();
+    }
+    return sum;
   }
 }