Added UpdateOneBusinessKeyTimestampStrategy

KAFKA-114

Author: nareshmaharaj-consultant

CHANGELOG.md

@@ -10,6 +10,7 @@
     Note: This defaults to false, which is a change of behaviour.
   - [KAFKA-118](https://jira.mongodb.org/browse/KAFKA-118) Made UuidStrategy configurable so can output BsonBinary Uuid values
   - [KAFKA-101](https://jira.mongodb.org/browse/KAFKA-101) Added `UuidProvidedInKeyStrategy` & `UuidProvidedInValueStrategy`
+  - [KAFKA-114](https://jira.mongodb.org/browse/KAFKA-114) Added `UpdateOneBusinessKeyTimestampStrategy` write model strategy`
 
 ## 1.1.0
   - [KAFKA-45](https://jira.mongodb.org/browse/KAFKA-45) Allow the Sink connector to ignore unused source record key or value fields.

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

@@ -0,0 +1,64 @@
+/*
+ * 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.MongoSinkTopicConfig.ID_FIELD;
+
+import java.time.Instant;
+
+import org.apache.kafka.connect.errors.DataException;
+
+import org.bson.BSONException;
+import org.bson.BsonDateTime;
+import org.bson.BsonDocument;
+
+import com.mongodb.client.model.UpdateOneModel;
+import com.mongodb.client.model.UpdateOptions;
+import com.mongodb.client.model.WriteModel;
+
+import com.mongodb.kafka.connect.sink.converter.SinkDocument;
+
+public class UpdateOneBusinessKeyTimestampStrategy implements WriteModelStrategy {
+
+    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";
+
+    @Override
+    public WriteModel<BsonDocument> createWriteModel(final SinkDocument document) {
+        BsonDocument vd = document.getValueDoc().orElseThrow(
+                () -> new DataException("Error: cannot build the WriteModel since the value document was missing unexpectedly"));
+
+        BsonDateTime dateTime = new BsonDateTime(Instant.now().toEpochMilli());
+
+        try {
+            BsonDocument businessKey = vd.getDocument(ID_FIELD);
+            vd.remove(ID_FIELD);
+
+            return new UpdateOneModel<>(
+                    businessKey,
+                    new BsonDocument("$set", vd.append(FIELD_NAME_MODIFIED_TS, dateTime))
+                            .append("$setOnInsert", new BsonDocument(FIELD_NAME_INSERTED_TS, dateTime)),
+                    UPDATE_OPTIONS);
+
+        } catch (BSONException e) {
+            throw new DataException("Error: cannot build the WriteModel since the value document does not contain an _id field of"
+                + " type BsonDocument which holds the business key fields");
+        }
+    }
+
+}

src/test/java/com/mongodb/kafka/connect/sink/MongoSinkConfigTest.java

@@ -96,6 +96,7 @@
 import com.mongodb.kafka.connect.sink.writemodel.strategy.DeleteOneDefaultStrategy;
 import com.mongodb.kafka.connect.sink.writemodel.strategy.ReplaceOneBusinessKeyStrategy;
 import com.mongodb.kafka.connect.sink.writemodel.strategy.ReplaceOneDefaultStrategy;
+import com.mongodb.kafka.connect.sink.writemodel.strategy.UpdateOneBusinessKeyTimestampStrategy;
 import com.mongodb.kafka.connect.sink.writemodel.strategy.UpdateOneTimestampsStrategy;
 import com.mongodb.kafka.connect.sink.writemodel.strategy.WriteModelStrategy;
 
@@ -415,6 +416,7 @@ Collection<DynamicTest> testGetSingleValidWriteModelStrategy() {
             put(ReplaceOneBusinessKeyStrategy.class.getName(), ReplaceOneBusinessKeyStrategy.class);
             put(ReplaceOneDefaultStrategy.class.getName(), ReplaceOneDefaultStrategy.class);
             put(UpdateOneTimestampsStrategy.class.getName(), UpdateOneTimestampsStrategy.class);
+            put(UpdateOneBusinessKeyTimestampStrategy.class.getName(), UpdateOneBusinessKeyTimestampStrategy.class);
         }};
 
         candidates.forEach((key, value) -> {
@@ -442,6 +444,7 @@ void testGetMultipleCollectionSpecificValidWriteModelStrategy() {
             put("topic-1", ReplaceOneDefaultStrategy.class);
             put("topic-2", ReplaceOneBusinessKeyStrategy.class);
             put("topic-3", UpdateOneTimestampsStrategy.class);
+            put("topic-5", UpdateOneBusinessKeyTimestampStrategy.class);
             put("topic-4", DeleteOneDefaultStrategy.class);
         }};
 

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

@@ -46,6 +46,9 @@ class WriteModelStrategyTest {
     private static final ReplaceOneBusinessKeyStrategy REPLACE_ONE_BUSINESS_KEY_STRATEGY = new ReplaceOneBusinessKeyStrategy();
     private static final UpdateOneTimestampsStrategy UPDATE_ONE_TIMESTAMPS_STRATEGY = new UpdateOneTimestampsStrategy();
 
+    private static final UpdateOneBusinessKeyTimestampStrategy
+        UPDATE_ONE_BUSINESS_KEY_TIMESTAMPS_STRATEGY = new UpdateOneBusinessKeyTimestampStrategy();
+
     private static final BsonDocument FILTER_DOC_DELETE_DEFAULT = BsonDocument.parse("{_id: {id: 1234}}");
     private static final BsonDocument FILTER_DOC_REPLACE_DEFAULT = BsonDocument.parse("{_id: 1234}");
     private static final BsonDocument REPLACEMENT_DOC_DEFAULT = BsonDocument.parse("{_id: 1234, first_name: 'Grace', last_name: 'Hopper'}");
@@ -189,4 +192,38 @@ void testUpdateOneTimestampsStrategyWithValidSinkDocument() {
         assertEquals(FILTER_DOC_UPDATE_TIMESTAMPS, writeModel.getFilter());
         assertTrue(writeModel.getOptions().isUpsert(), "update expected to be done in upsert mode");
     }
+
+    @Test
+    @DisplayName("when value document is missing for UpdateOneBusinessKeyTimestampStrategy then DataException")
+    void testUpdateOneBusinessKeyTimestampsStrategyWithMissingValueDocument() {
+        assertThrows(DataException.class,
+            () -> UPDATE_ONE_BUSINESS_KEY_TIMESTAMPS_STRATEGY.createWriteModel(new SinkDocument(new BsonDocument(), null)));
+    }
+
+    @Test
+    @DisplayName("when sink document is valid for UpdateOneBusinessKeyTimestampStrategy then correct UpdateOneModel")
+    void testUpdateOneBusinessKeyTimestampsStrategyWithValidSinkDocument() {
+        BsonDocument valueDoc = BsonDocument.parse("{_id: 1234, first_name: 'Grace', last_name: 'Hopper', active: false}");
+
+        WriteModel<BsonDocument> result = UPDATE_ONE_TIMESTAMPS_STRATEGY.createWriteModel(new SinkDocument(null, valueDoc));
+        assertTrue(result instanceof UpdateOneModel, "result expected to be of type UpdateOneModel");
+
+        UpdateOneModel<BsonDocument> writeModel = (UpdateOneModel<BsonDocument>) result;
+
+        //NOTE: This test case can only check:
+        // i) for both fields to be available
+        // ii) having the correct BSON type (BsonDateTime)
+        // iii) and be initially equal
+        // The exact dateTime value is not directly testable here.
+        BsonDocument updateDoc = (BsonDocument) writeModel.getUpdate();
+
+        BsonDateTime modifiedTS = updateDoc.getDocument("$set").getDateTime(UpdateOneBusinessKeyTimestampStrategy.FIELD_NAME_MODIFIED_TS);
+        BsonDateTime insertedTS = updateDoc.getDocument("$setOnInsert")
+            .getDateTime(UpdateOneBusinessKeyTimestampStrategy.FIELD_NAME_INSERTED_TS);
+
+        assertEquals(insertedTS, modifiedTS, "modified and inserted timestamps must initially be equal");
+        assertTrue(writeModel.getFilter() instanceof BsonDocument, "filter expected to be of type BsonDocument");
+        assertEquals(FILTER_DOC_UPDATE_TIMESTAMPS, writeModel.getFilter());
+        assertTrue(writeModel.getOptions().isUpsert(), "update expected to be done in upsert mode");
+    }
 }