feat: support iterable for druid lookups-cached-single extension

Signed-off-by: TessaIO <[email protected]>
apache · Mar 14, 2024 · 8e417d2 · 8e417d2
1 parent b8201e3
commit 8e417d2
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 9 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
@@ -31,6 +31,7 @@
 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;
@@ -109,25 +110,32 @@ public List<String> unapply(@Nullable final String value)
   @Override
   public boolean canIterate()
   {
-    return false;
+    return true;
   }
 
   @Override
   public boolean canGetKeySet()
   {
-    return false;
+    return true;
   }
 
   @Override
   public Iterable<Map.Entry<String, String>> iterable()
   {
-    throw new UnsupportedOperationException("Cannot iterate");
+    return this.dataFetcher.fetchAll();
   }
 
   @Override
   public Set<String> keySet()
   {
-    throw new UnsupportedOperationException("Cannot get key set");
+    Iterable<Map.Entry<String, String>> data = this.dataFetcher.fetchAll();
+    Set<String> set = new HashSet<>();
+    if (data != null) {
+      data.forEach(each -> {
+        set.add(each.getKey());
+      });
+    }
+    return set;
   }
 
   @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
@@ -19,6 +19,7 @@
 
 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;
@@ -30,8 +31,7 @@
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 
-import java.util.Arrays;
-import java.util.Collections;
+import java.util.*;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 
@@ -132,13 +132,38 @@ public void testGetCacheKey()
   @Test
   public void testCanGetKeySet()
   {
-    Assert.assertFalse(loadingLookup.canGetKeySet());
+    Assert.assertTrue(loadingLookup.canGetKeySet());
+  }
+
+  @Test
+  public void testCanIterate()
+  {
+    Assert.assertTrue(loadingLookup.canIterate());
   }
 
   @Test
   public void testKeySet()
   {
-    expectedException.expect(UnsupportedOperationException.class);
-    loadingLookup.keySet();
+    Map.Entry<String, String> fetchedData = new AbstractMap.SimpleEntry<>("dummy", "test");
+    EasyMock.expect(dataFetcher.fetchAll()).andReturn(Arrays.asList(fetchedData));
+    EasyMock.replay(dataFetcher);
+    Assert.assertEquals(loadingLookup.keySet().size(), 1);
+    EasyMock.verify(dataFetcher);
+  }
+
+  @Test
+  public void testFetchAll()
+  {
+    Map.Entry<String, String> fetchedData = new AbstractMap.SimpleEntry<>("dummy", "test");
+    EasyMock.expect(dataFetcher.fetchAll()).andReturn(Arrays.asList(fetchedData));
+    EasyMock.replay(dataFetcher);
+    Assert.assertEquals(getIteratorSize(loadingLookup.iterable().iterator()), 1);
+    EasyMock.verify(dataFetcher);
+  }
+
+  public int getIteratorSize(Iterator<Map.Entry<String, String>> it) {
+    int i = 0;
+    for ( ; it.hasNext() ; ++i ) it.next();
+    return i;
   }
 }