Fix business key update strategies to use dot notation for filters (#39)

KAFKA-155

Author: rozza

CHANGELOG.md

@@ -20,6 +20,7 @@
   - [KAFKA-154](https://jira.mongodb.org/browse/KAFKA-154) Improve the handling and error messaging for Json array config values.
   - [KAFKA-78](https://jira.mongodb.org/browse/KAFKA-78) Added dead letter queue support for the source connector.
   - [KAFKA-157](https://jira.mongodb.org/browse/KAFKA-157) Improved error message for business key errors.
+  - [KAFKA-155](https://jira.mongodb.org/browse/KAFKA-155) Fix business key update strategies to use dot notation for filters
 
 
 ## 1.2.0

src/main/java/com/mongodb/kafka/connect/sink/writemodel/strategy/ReplaceOneBusinessKeyStrategy.java

@@ -19,6 +19,7 @@
 package com.mongodb.kafka.connect.sink.writemodel.strategy;
 
 import static com.mongodb.kafka.connect.sink.MongoSinkTopicConfig.ID_FIELD;
+import static com.mongodb.kafka.connect.sink.writemodel.strategy.WriteModelHelper.flattenKeys;
 
 import org.apache.kafka.connect.errors.DataException;
 
@@ -29,11 +30,17 @@
 import com.mongodb.client.model.ReplaceOptions;
 import com.mongodb.client.model.WriteModel;
 
+import com.mongodb.kafka.connect.sink.Configurable;
+import com.mongodb.kafka.connect.sink.MongoSinkTopicConfig;
 import com.mongodb.kafka.connect.sink.converter.SinkDocument;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.IdStrategy;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.PartialKeyStrategy;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.PartialValueStrategy;
 
-public class ReplaceOneBusinessKeyStrategy implements WriteModelStrategy {
+public class ReplaceOneBusinessKeyStrategy implements WriteModelStrategy, Configurable {
 
   private static final ReplaceOptions REPLACE_OPTIONS = new ReplaceOptions().upsert(true);
+  private boolean isPartialId = false;
 
   @Override
   public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
@@ -48,6 +55,9 @@ public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
     try {
       BsonDocument businessKey = vd.getDocument(ID_FIELD);
       vd.remove(ID_FIELD);
+      if (isPartialId) {
+        businessKey = flattenKeys(businessKey);
+      }
       return new ReplaceOneModel<>(businessKey, vd, REPLACE_OPTIONS);
     } catch (BSONException e) {
       throw new DataException(
@@ -56,4 +66,11 @@ public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
               + " existing `_id` value in the business key then ensure `document.id.strategy.overwrite.existing=true`.");
     }
   }
+
+  @Override
+  public void configure(final MongoSinkTopicConfig configuration) {
+    IdStrategy idStrategy = configuration.getIdStrategy();
+    isPartialId =
+        idStrategy instanceof PartialKeyStrategy || idStrategy instanceof PartialValueStrategy;
+  }
 }

src/main/java/com/mongodb/kafka/connect/sink/writemodel/strategy/UpdateOneBusinessKeyTimestampStrategy.java

@@ -17,6 +17,7 @@
 package com.mongodb.kafka.connect.sink.writemodel.strategy;
 
 import static com.mongodb.kafka.connect.sink.MongoSinkTopicConfig.ID_FIELD;
+import static com.mongodb.kafka.connect.sink.writemodel.strategy.WriteModelHelper.flattenKeys;
 
 import java.time.Instant;
 
@@ -30,13 +31,19 @@
 import com.mongodb.client.model.UpdateOptions;
 import com.mongodb.client.model.WriteModel;
 
+import com.mongodb.kafka.connect.sink.Configurable;
+import com.mongodb.kafka.connect.sink.MongoSinkTopicConfig;
 import com.mongodb.kafka.connect.sink.converter.SinkDocument;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.IdStrategy;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.PartialKeyStrategy;
+import com.mongodb.kafka.connect.sink.processor.id.strategy.PartialValueStrategy;
 
-public class UpdateOneBusinessKeyTimestampStrategy implements WriteModelStrategy {
+public class UpdateOneBusinessKeyTimestampStrategy implements WriteModelStrategy, Configurable {
 
   private static final UpdateOptions UPDATE_OPTIONS = new UpdateOptions().upsert(true);
   static final String FIELD_NAME_MODIFIED_TS = "_modifiedTS";
   static final String FIELD_NAME_INSERTED_TS = "_insertedTS";
+  private boolean isPartialId = false;
 
   @Override
   public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
@@ -53,7 +60,9 @@ public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
     try {
       BsonDocument businessKey = vd.getDocument(ID_FIELD);
       vd.remove(ID_FIELD);
-
+      if (isPartialId) {
+        businessKey = flattenKeys(businessKey);
+      }
       return new UpdateOneModel<>(
           businessKey,
           new BsonDocument("$set", vd.append(FIELD_NAME_MODIFIED_TS, dateTime))
@@ -67,4 +76,11 @@ public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
               + "existing `_id` value in the business key then ensure `document.id.strategy.overwrite.existing=true`.");
     }
   }
+
+  @Override
+  public void configure(final MongoSinkTopicConfig configuration) {
+    IdStrategy idStrategy = configuration.getIdStrategy();
+    isPartialId =
+        idStrategy instanceof PartialKeyStrategy || idStrategy instanceof PartialValueStrategy;
+  }
 }

src/main/java/com/mongodb/kafka/connect/sink/writemodel/strategy/WriteModelHelper.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.kafka.connect.sink.writemodel.strategy;
+
+import static java.lang.String.format;
+import static java.util.Collections.singletonList;
+import static java.util.stream.Collectors.toList;
+
+import java.util.List;
+
+import org.bson.BsonDocument;
+import org.bson.BsonElement;
+
+final class WriteModelHelper {
+
+  private static final String CREATE_PREFIX = "%s%s.";
+  private static final String ELEMENT_NAME_PREFIX = "%s%s";
+
+  static BsonDocument flattenKeys(final BsonDocument original) {
+    BsonDocument businessKey = new BsonDocument();
+    original.forEach(
+        (key, value) ->
+            flattenBsonElement("", new BsonElement(key, value))
+                .forEach(b -> businessKey.append(b.getName(), b.getValue())));
+    return businessKey;
+  }
+
+  static List<BsonElement> flattenBsonElement(final String prefix, final BsonElement element) {
+    if (element.getValue().isDocument()) {
+      return element.getValue().asDocument().entrySet().stream()
+          .flatMap(
+              e ->
+                  flattenBsonElement(
+                      format(CREATE_PREFIX, prefix, element.getName()),
+                      new BsonElement(e.getKey(), e.getValue()))
+                      .stream())
+          .collect(toList());
+    }
+    return singletonList(
+        new BsonElement(
+            format(ELEMENT_NAME_PREFIX, prefix, element.getName()), element.getValue()));
+  }
+
+  private WriteModelHelper() {}
+}

src/test/java/com/mongodb/kafka/connect/sink/writemodel/strategy/WriteModelHelperTest.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2008-present MongoDB, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.mongodb.kafka.connect.sink.writemodel.strategy;
+
+import static com.mongodb.kafka.connect.sink.writemodel.strategy.WriteModelHelper.flattenKeys;
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+import org.junit.platform.runner.JUnitPlatform;
+import org.junit.runner.RunWith;
+
+import org.bson.BsonDocument;
+
+@RunWith(JUnitPlatform.class)
+class WriteModelHelperTest {
+
+  @Test
+  @DisplayName("flatten keys works as expected")
+  void testFlattenKeys() {
+    assertAll(
+        () -> assertEquals(BsonDocument.parse("{}"), flattenKeys(BsonDocument.parse("{}"))),
+        () ->
+            assertEquals(
+                BsonDocument.parse("{a: 1, b: 1, c: [1]}"),
+                flattenKeys(BsonDocument.parse("{a: 1, b: 1, c: [1]}"))),
+        () ->
+            assertEquals(
+                BsonDocument.parse(
+                    "{'a.a': 1, 'b.b': 1, 'b.c.d': 1, 'b.c.e.f': 1, g: [{h: 1, i: {j: 1}}]}"),
+                flattenKeys(
+                    BsonDocument.parse(
+                        "{a: {a: 1}, b: {b: 1, c: {d: 1, e: {f: 1}}}, g: [{h: 1, i: {j: 1}}]}"))));
+  }
+}