From 8e417d2b012661884b33d61731f346feb7ac35f2 Mon Sep 17 00:00:00 2001
From: TessaIO <ahmedgrati1999@gmail.com>
Date: Thu, 14 Mar 2024 16:31:02 +0100
Subject: [PATCH] feat: support iterable for druid lookups-cached-single
 extension

Signed-off-by: TessaIO <ahmedgrati1999@gmail.com>
---
 .../druid/server/lookup/LoadingLookup.java    | 16 ++++++---
 .../server/lookup/LoadingLookupTest.java      | 35 ++++++++++++++++---
 2 files changed, 42 insertions(+), 9 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 2bffc36ee808..58cbfb4e32e2 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
@@ -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/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 0a28454aa58c..4360b9178691 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
@@ -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;
   }
 }