Handle empty values for task and datasource conditions in pod templat…

…e selector (apache#17400) * handling empty sets for dataSourceCondition and taskTypeCondition * using new HashSet<>() to fix forbidden api error in testCheck * fixing style issues
miqdigital · Oct 31, 2024 · 5fcf420 · 5fcf420
1 parent 4b7902e
commit 5fcf420
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 3 deletions.
diff --git a/...s-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/execution/Selector.java b/...s-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/execution/Selector.java
@@ -23,6 +23,7 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.indexing.common.task.Task;
 import org.apache.druid.query.DruidMetrics;
+import org.apache.druid.utils.CollectionUtils;
 
 import java.util.Map;
 import java.util.Objects;
@@ -70,6 +71,7 @@ public Selector(
   public boolean evaluate(Task task)
   {
     boolean isMatch = true;
+
     if (cxtTagsConditions != null) {
       isMatch = cxtTagsConditions.entrySet().stream().allMatch(entry -> {
         String tagKey = entry.getKey();
@@ -80,15 +82,15 @@ public boolean evaluate(Task task)
         }
         Object tagValue = tags.get(tagKey);
 
-        return tagValue == null ? false : tagValues.contains((String) tagValue);
+        return tagValue != null && tagValues.contains((String) tagValue);
       });
     }
 
-    if (isMatch && taskTypeCondition != null) {
+    if (isMatch && !CollectionUtils.isNullOrEmpty(taskTypeCondition)) {
       isMatch = taskTypeCondition.contains(task.getType());
     }
 
-    if (isMatch && dataSourceCondition != null) {
+    if (isMatch && !CollectionUtils.isNullOrEmpty(dataSourceCondition)) {
       isMatch = dataSourceCondition.contains(task.getDataSource());
     }
 

diff --git a/...erlord-extensions/src/test/java/org/apache/druid/k8s/overlord/execution/SelectorTest.java b/...erlord-extensions/src/test/java/org/apache/druid/k8s/overlord/execution/SelectorTest.java
@@ -30,11 +30,68 @@
 import org.junit.Test;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 
 public class SelectorTest
 {
+  @Test
+  public void shouldReturnTrueWhenMatchTasksTagsAndEmptyDataSource()
+  {
+    Map<String, Set<String>> cxtTagsConditions = new HashMap<>();
+    cxtTagsConditions.put("tag1", Sets.newHashSet("tag1Value"));
+
+    Task task = NoopTask.create();
+    task.addToContext(DruidMetrics.TAGS, ImmutableMap.of("tag1", "tag1Value"));
+
+    Selector selector = new Selector(
+        "TestSelector",
+        cxtTagsConditions,
+        Sets.newHashSet(NoopTask.TYPE),
+        new HashSet<>()
+    );
+
+    Assert.assertTrue(selector.evaluate(task));
+  }
+
+  @Test
+  public void shouldReturnTrueWhenMatchDataSourceTagsAndEmptyTasks()
+  {
+    String datasource = "table";
+    Map<String, Set<String>> cxtTagsConditions = new HashMap<>();
+    cxtTagsConditions.put("tag1", Sets.newHashSet("tag1Value"));
+
+    Task task = NoopTask.forDatasource(datasource);
+    task.addToContext(DruidMetrics.TAGS, ImmutableMap.of("tag1", "tag1Value"));
+
+    Selector selector = new Selector(
+        "TestSelector",
+        cxtTagsConditions,
+        new HashSet<>(),
+        Sets.newHashSet(datasource)
+    );
+
+    Assert.assertTrue(selector.evaluate(task));
+  }
+
+  @Test
+  public void shouldReturnTrueWhenMatchDataSourceTasksAndEmptyTags()
+  {
+    String datasource = "table";
+    Map<String, Set<String>> cxtTagsConditions = new HashMap<>();
+
+    Task task = NoopTask.forDatasource(datasource);
+
+    Selector selector = new Selector(
+        "TestSelector",
+        cxtTagsConditions,
+        Sets.newHashSet(NoopTask.TYPE),
+        Sets.newHashSet(datasource)
+    );
+
+    Assert.assertTrue(selector.evaluate(task));
+  }
 
   @Test
   public void shouldReturnTrueWhenAllTagsAndTasksMatch()