Repository: microsoft/kafka-connect-cosmosdb (java)

Issue:
Set the message key in Kafka topics
**## Problem Statement**
Currently message key in Kafka is not set by the Source connector

**## Proposed Solution**
Use **ExtractField** & **ValueToKey** SMT to set the message key to the value of the _id_ field from the Cosmos DB document

**## Acceptance Criteria**

- [ ] If the approach in the comments works: An integration test and/or walkthrough/tutorial document using SMT to achieve this functionality showing the correct set of configurations.
- [ ] If the approach does not work: Create a new work item with findings and information.

**## Additional Context**
Refer to comments in the discussion for usage and examples of using **ExtractField** and **ValueToKey**

Pull request:
Adding option to set message key
## Type of PR

- [x] Documentation changes
- [x] Code changes
- [ ] Test changes
- [ ] CI-CD changes
- [ ] GitHub Template changes

## Purpose of PR
- Setting the message key using the Cosmos DB item ID
- Added a source setting to enable/disable this message key behavior
- Updated docs for info on new source setting

## Issues Closed or Referenced

- Closes #159 

Patch:
diff --git a/doc/README_Source.md b/doc/README_Source.md
index c0ae607..75b1550 100644
--- a/doc/README_Source.md
+++ b/doc/README_Source.md
@@ -154,7 +154,6 @@ To delete the created Azure Cosmos DB service and its resource group using Azure
 
 The following settings are used to configure the Cosmos DB Kafka Source Connector. These configuration values determine which Cosmos DB container is consumed, which Kafka topics data is written into and formats to serialize the data. For an example configuration file with the default values, refer to [this config](../src/integration-test/resources/source.config.json).
 
-
 | Name | Type | Description | Required/Optional |
 |------|------|-------------|-------------------|
 | connector.class | string | Classname of the Cosmos DB sink. Should be set to `com.microsoft.azure.cosmosdb.kafka.connect.sink.CosmosDBSourceConnector` | Required |
@@ -163,10 +162,12 @@ The following settings are used to configure the Cosmos DB Kafka Source Connecto
 | connect.cosmosdb.connection.endpoint | uri | the endpoint for the Cosmos DB Account | Required |
 | connect.cosmosdb.containers.topicmap | string | comma separeted topic to collection mapping, eg. topic1#coll1,topic2#coll2 | Required |
 | connect.cosmosdb.containers | string | list of collections to monitor |  Required |
-| connect.cosmosdb.changefeed.startFromBeginning | boolean |  set if the change feed should start from beginning | Required |
-| connect.cosmosdb.task.poll.interval | int | interval to poll the changefeedcontainer for changes | Required | 
+| connect.cosmosdb.messagekey.enabled | boolean | set if the Kafka message key should be set. Default is `true` | Required |
+| connect.cosmosdb.messagekey.field | string | use the field's value from the document as the message key. Default is `id` | Required |
+| connect.cosmosdb.changefeed.startFromBeginning | boolean |  set if the change feed should start from beginning. Default is `true` | Required |
+| connect.cosmosdb.task.poll.interval | int | interval to poll the changefeedcontainer for changes | Required |
 | key.converter | string | Serialization format for the key data written into Kafka topic | Required |
 | value.converter | string | Serialization format for the value data written into the Kafka topic | Required |
 | key.converter.schemas.enable | string | Set to `"true"` if the key data has embedded schema | Optional |
 | value.converter.schemas.enable | string | Set to `"true"` if the key data has embedded schema | Optional |
-| tasks.max | int | Maximum number of connector sink tasks. Default is `1` | Optional |
\ No newline at end of file
+| tasks.max | int | Maximum number of connector sink tasks. Default is `1` | Optional |
diff --git a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTask.java b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTask.java
index 3ab5675..c130e51 100644
--- a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTask.java
+++ b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTask.java
@@ -4,6 +4,7 @@ import com.azure.cosmos.*;
 import com.azure.cosmos.models.*;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.microsoft.azure.cosmosdb.kafka.connect.TopicContainerMap;
+import org.apache.kafka.connect.data.Schema;
 import org.apache.kafka.connect.source.SourceRecord;
 import org.apache.kafka.connect.source.SourceTask;
 import org.slf4j.Logger;
@@ -90,8 +91,18 @@ public class CosmosDBSourceTask extends SourceTask {
                 if(node == null) {
                     continue;
                 }
+
+                // Set the Kafka message key if option is enabled and field is configured in document
+                String messageKey = "";
+                if (this.settings.isMessageKeyEnabled()) {
+                    JsonNode messageKeyFieldNode = node.get(this.settings.getMessageKeyField());
+                    messageKey = (messageKeyFieldNode != null) ? messageKeyFieldNode.toString() : "";
+                }
+
                 // Since Lease container takes care of maintaining state we don't have to send source offset to kafka
-                SourceRecord sourceRecord = new SourceRecord(partition, null, topic, null, node.toString());
+                SourceRecord sourceRecord = new SourceRecord(partition, null, topic,
+                                                    Schema.STRING_SCHEMA, messageKey, 
+                                                    Schema.STRING_SCHEMA, node.toString());
                 bufferSize -= sourceRecord.value().toString().getBytes().length;
                 // If the buffer Size exceeds then do not remove the node .
                 if (bufferSize <=0){
diff --git a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettingDefaults.java b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettingDefaults.java
index afa6cd7..9407cba 100644
--- a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettingDefaults.java
+++ b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettingDefaults.java
@@ -2,5 +2,7 @@ package com.microsoft.azure.cosmosdb.kafka.connect.source;
 
 public class SourceSettingDefaults {
 
+    public static final Boolean MESSAGE_KEY_ENABLED = true;
+    public static final String MESSAGE_KEY_FIELD = "id";
     public static final Boolean CHANGE_FEED_START_FROM_BEGINNING = true;
 }
diff --git a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettings.java b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettings.java
index db5b645..58a52aa 100644
--- a/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettings.java
+++ b/src/main/java/com/microsoft/azure/cosmosdb/kafka/connect/source/SourceSettings.java
@@ -14,6 +14,7 @@ import java.util.List;
  * Contains settings for the CosmosDB Kafka Source Connector
  */
 public class SourceSettings extends Settings {
+    private boolean messageKeyEnabled;
     private boolean startFromBeginning;
 
     private final List<Setting> sourceSettings = Arrays.asList(
@@ -32,6 +33,10 @@ public class SourceSettings extends Settings {
                     "Assigned Container", this::setAssignedContainer, this::getAssignedContainer),
             new Setting(Settings.PREFIX + ".worker.name", "The Cosmos DB worker name.",
                     "Worker name", this::setWorkerName, this::getWorkerName),
+            new BooleanSetting(Settings.PREFIX + ".messagekey.enabled", "Whether to set the Kafka message key.",
+                    "Message key enabled", SourceSettingDefaults.MESSAGE_KEY_ENABLED, this::setMessageKeyEnabled, this::isMessageKeyEnabled),
+            new Setting(Settings.PREFIX + ".messagekey.field", "The document field to use as the message key.",
+                    "Message key field", SourceSettingDefaults.MESSAGE_KEY_FIELD, this::setMessageKeyField, this::getMessageKeyField),
             new BooleanSetting(Settings.PREFIX + ".changefeed.startFromBeginning", "Whether the change feed should start from beginning.",
                     "Change Feed start from beginning", SourceSettingDefaults.CHANGE_FEED_START_FROM_BEGINNING, this::setStartFromBeginning, this::isStartFromBeginning)
     );
@@ -57,6 +62,20 @@ public class SourceSettings extends Settings {
         this.assignedContainer = assignedPartitions;
     }
 
+    public boolean isMessageKeyEnabled() { return this.messageKeyEnabled; }
+
+    public void setMessageKeyEnabled(boolean messageKeyEnabled) { this.messageKeyEnabled = messageKeyEnabled;  }
+
+    private String messageKeyField;
+
+    public String getMessageKeyField() {
+        return this.messageKeyField;
+    }
+
+    public void setMessageKeyField(String messageKeyField) {
+        this.messageKeyField = messageKeyField;
+    }
+
     public boolean isStartFromBeginning() { return this.startFromBeginning; }
 
     public void setStartFromBeginning(boolean startFromBeginning) { this.startFromBeginning = startFromBeginning;  }


Test patch:
diff --git a/src/test/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTaskTest.java b/src/test/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTaskTest.java
index 83244fb..11a6c1a 100644
--- a/src/test/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTaskTest.java
+++ b/src/test/java/com/microsoft/azure/cosmosdb/kafka/connect/source/CosmosDBSourceTaskTest.java
@@ -90,6 +90,25 @@ public class CosmosDBSourceTaskTest {
         Assert.assertEquals(result.size(),1);
     }
 
+    @Test
+    public void testPollWithMessageKey() throws InterruptedException, JsonProcessingException {
+        String jsonString = "{\"id\":123,\"k1\":\"v1\",\"k2\":\"v2\"}";
+        ObjectMapper mapper = new ObjectMapper();
+        JsonNode actualObj = mapper.readTree(jsonString);
+        List<JsonNode> changes = new ArrayList<>();
+        settings.setMessageKeyEnabled(true);
+        settings.setMessageKeyField("id");
+        changes.add(actualObj);
+
+        new Thread(() -> {
+            testTask.handleCosmosDbChanges(changes);
+        }).start();
+
+        List<SourceRecord>  result=testTask.poll();
+        Assert.assertEquals(result.size(),1);
+        Assert.assertEquals(result.get(0).key(), "123");
+    }
+
     @Test
     public void testZeroBatchSize() throws InterruptedException, JsonProcessingException, IllegalAccessException {
         String jsonString = "{\"k1\":\"v1\",\"k2\":\"v2\"}";