Merge pull request #49 from johnmbw/DSTORE-1475-make-conditions-suppo…

…rt-atomic-updates DSTORE-1475 handle atomic updates for conditions
BrandwatchLtd · Apr 19, 2021 · 18a54ba · 18a54ba
2 parents a6b254f + 1d4526d
commit 18a54ba
Show file tree

Hide file tree

Showing 2 changed files with 108 additions and 1 deletion.
diff --git a/solr/core/src/java/org/apache/solr/update/processor/UpsertCondition.java b/solr/core/src/java/org/apache/solr/update/processor/UpsertCondition.java
@@ -29,6 +29,7 @@
 import java.util.function.Predicate;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
+import java.util.stream.Stream;
 
 import org.apache.lucene.search.BooleanClause;
 import org.apache.solr.common.SolrException;
@@ -249,12 +250,35 @@ private static Predicate<Object> stringlyEquals(String value) {
       };
     }
 
+    private static Predicate<Map<?,?>> forAtomicUpdate(Predicate<Object> fieldPredicate) {
+      return fieldValue -> {
+        if (fieldValue.containsKey("set") || fieldValue.containsKey("add")) {
+          return fieldValue.values().stream()
+              .flatMap(updateValue -> {
+                if (updateValue instanceof Collection) {
+                  return ((Collection<?>)updateValue).stream();
+                }
+                return Stream.of(updateValue);
+              })
+              .anyMatch(fieldPredicate);
+        }
+        return false;
+      };
+    }
+
     private static Predicate<SolrInputDocument> forField(String field, Predicate<Object> fieldPredicate) {
+      Predicate<Map<?,?>> atomicUpdatePredicate = forAtomicUpdate(fieldPredicate);
+      Predicate<Object> predicate = fieldValue -> {
+        if (fieldValue instanceof Map) {
+          return atomicUpdatePredicate.test((Map<?,?>)fieldValue);
+        }
+        return fieldPredicate.test(fieldValue);
+      };
       return doc -> {
         if (doc != null) {
           Collection<Object> values = doc.getFieldValues(field);
           if (values != null) {
-            return values.stream().anyMatch(fieldPredicate);
+            return values.stream().anyMatch(predicate);
           }
         }
         return false;

diff --git a/solr/core/src/test/org/apache/solr/update/processor/UpsertConditionTest.java b/solr/core/src/test/org/apache/solr/update/processor/UpsertConditionTest.java
@@ -1,8 +1,10 @@
 package org.apache.solr.update.processor;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ListMultimap;
@@ -178,6 +180,87 @@ public void givenNumericField_whenMatching() {
     assertTrue(condition.matches(oldDoc, newDoc));
   }
 
+  @Test
+  public void givenAtomicUpdateSet_whenMatching() {
+    NamedList<String> args = new NamedList<>(ImmutableMap.of(
+        "must", "NEW.field:value",
+        "action", "skip"
+    ));
+
+    UpsertCondition condition = UpsertCondition.parse("skip-it", args);
+
+    assertThat(condition.getName(), is("skip-it"));
+
+    SolrInputDocument oldDoc = new SolrInputDocument();
+    SolrInputDocument newDoc = new SolrInputDocument();
+
+    newDoc.setField("field", Collections.singletonMap("set", "other1"));
+    assertFalse(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("set", "value"));
+    assertTrue(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("set", ImmutableList.of("value", "other2")));
+    assertTrue(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("set", ImmutableList.of("other1", "other2")));
+    assertFalse(condition.matches(oldDoc, newDoc));
+  }
+
+  @Test
+  public void givenAtomicUpdateAdd_whenMatching() {
+    NamedList<String> args = new NamedList<>(ImmutableMap.of(
+        "must", "NEW.field:value",
+        "action", "skip"
+    ));
+
+    UpsertCondition condition = UpsertCondition.parse("skip-it", args);
+
+    assertThat(condition.getName(), is("skip-it"));
+
+    SolrInputDocument oldDoc = new SolrInputDocument();
+    SolrInputDocument newDoc = new SolrInputDocument();
+
+    newDoc.setField("field", Collections.singletonMap("add", "other1"));
+    assertFalse(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("add", "value"));
+    assertTrue(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("add", ImmutableList.of("value", "other2")));
+    assertTrue(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("add", ImmutableList.of("other1", "other2")));
+    assertFalse(condition.matches(oldDoc, newDoc));
+  }
+
+  @Test
+  public void givenAtomicUpdateRemove_whenMatching() {
+    NamedList<String> args = new NamedList<>(ImmutableMap.of(
+        "must", "NEW.field:value",
+        "action", "skip"
+    ));
+
+    UpsertCondition condition = UpsertCondition.parse("skip-it", args);
+
+    assertThat(condition.getName(), is("skip-it"));
+
+    SolrInputDocument oldDoc = new SolrInputDocument();
+    SolrInputDocument newDoc = new SolrInputDocument();
+
+    oldDoc.setField("field", Collections.singletonMap("remove", "other1"));
+    assertFalse(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("remove", "value"));
+    assertFalse(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("remove", ImmutableList.of("value", "other2")));
+    assertFalse(condition.matches(oldDoc, newDoc));
+
+    newDoc.setField("field", Collections.singletonMap("remove", ImmutableList.of("other1", "other2")));
+    assertFalse(condition.matches(oldDoc, newDoc));
+  }
+
   @Test
   public void givenSingleMustClause_whenMatching() {
     NamedList<String> args = namedList(ImmutableListMultimap.of(