Repository: microsoft/kafka-connect-cosmosdb
Pull request #492 (2022-10-13 22:43:10)

Release1.6.0
Merge from dev to main for the 1.6.0 release


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index d483610..e0acd28 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,12 +1,17 @@
 ## Release History
 
+### 1.6.0 (2022-10-13)
+#### New Features
+* Added bulk support for `CosmosDBSinkConnector`. - [PR 485](https://github.com/microsoft/kafka-connect-cosmosdb/pull/485)
+* Added gateway mode support for `CosmosDBSinkConnector` and `CosmosDBSourceConnector`. - [PR 490](https://github.com/microsoft/kafka-connect-cosmosdb/pull/490)
+
 ### 1.5.0 (2022-09-16)
 #### New Features
 * Updated `azure-cosmos` version to 4.36.0.
 
 #### Key Bug Fixes
 * Fixed parsing for empty arrays. [PR 466](https://github.com/microsoft/kafka-connect-cosmosdb/pull/466)
-* Updated `CosmosDBSinkConnector` to keep retrying for throttled requests. [PR 472](https://github.com/microsoft/kafka-connect-cosmosdb/pull/472)
+* Updated `CosmosDBSinkConnector` to keep retrying for throttled requests. - [PR 472](https://github.com/microsoft/kafka-connect-cosmosdb/pull/472)
 
 ### 1.4.0 (2022-05-26)
 #### New Features
diff --git a/pom.xml b/pom.xml
index 4937c82..b3dacd4 100644
--- a/pom.xml
+++ b/pom.xml
@@ -7,7 +7,7 @@
 
     <groupId>com.azure.cosmos.kafka</groupId>
     <artifactId>kafka-connect-cosmos</artifactId>
-    <version>1.5.0</version>
+    <version>1.6.0</version>
 
     <name> kafka-connect-cosmos</name>
     <url>https://github.com/microsoft/kafka-connect-cosmosdb</url>
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java b/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java
index 8afe94d..4576ac3 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect;
 
 import static com.azure.cosmos.kafka.connect.CosmosDBConfig.CosmosClientBuilder.createClient;
@@ -50,7 +53,27 @@ public class CosmosDBConfig extends AbstractConfig {
     private static final String COSMOS_CONTAINER_TOPIC_MAP_DISPLAY = "Topic-Container map";
     private static final String COSMOS_CONTAINER_TOPIC_MAP_DOC =
             "A comma delimited list of Kafka topics mapped to Cosmos containers.\n"
-                    + "For example: topic1#con1,topic2#con2.";    
+                    + "For example: topic1#con1,topic2#con2.";
+
+    public static final String COSMOS_GATEWAY_MODE_ENABLED = "connect.cosmos.connection.gateway.enabled";
+    private static final String COSMOS_GATEWAY_MODE_ENABLED_DOC =
+            "Flag to indicate whether to use gateway mode. By default it is false";
+    private static final boolean DEFAULT_COSMOS_GATEWAY_MODE_ENABLED = false;
+
+    public static final String COSMOS_CONNECTION_SHARING_ENABLED = "connect.cosmos.connection.sharing.enabled";
+    private static final String COSMOS_CONNECTION_SHARING_ENABLED_DOC =
+            "If you have set 'connect.cosmos.connection.gateway.enabled' to true, then this configure will not make any difference. "
+            + "By enabling this it allows connection sharing between instances of cosmos clients on the same jvm.";
+    private static final boolean DEFAULT_COSMOS_CONNECTION_SHARING_ENABLED = false;
+
+    public static final String COSMOS_SINK_BULK_ENABLED_CONF = "connect.cosmos.sink.bulk.enabled";
+    private static final String COSMOS_SINK_BULK_ENABLED_DOC = "Flag to indicate whether Cosmos DB bulk mode is enabled for Sink connector. By default it is true.";
+    private static final boolean DEFAULT_COSMOS_SINK_BULK_ENABLED = true;
+
+    public static final String COSMOS_SINK_MAX_RETRY_COUNT = "connect.cosmos.sink.maxRetryCount";
+    private static final String COSMOS_SINK_MAX_RETRY_COUNT_DOC =
+            "Cosmos DB max retry attempts on write failures for Sink connector. By default, the connector will retry on transient write errors for up to 10 times.";
+    private static final int DEFAULT_COSMOS_SINK_MAX_RETRY_COUNT = 10;
         
     public static final String COSMOS_PROVIDER_NAME_CONF = "connect.cosmos.provider.name";
     private static final String COSMOS_PROVIDER_NAME_DEFAULT = null;
@@ -63,20 +86,28 @@ public class CosmosDBConfig extends AbstractConfig {
     public static final int COSMOS_DATABASE_GROUP_ORDER = 2;
     public static final String COSMOS_CLIENT_USER_AGENT_SUFFIX = "APN/1.0 Microsoft/1.0 KafkaConnect/";
 
-    private String connEndpoint;
-    private String connKey;
-    private String databaseName;
-    private String providerName;
+    private final String connEndpoint;
+    private final String connKey;
+    private final String databaseName;
+    private final String providerName;
+    private final boolean bulkModeEnabled;
+    private final boolean gatewayModeEnabled;
+    private final boolean connectionSharingEnabled;
+    private final int maxRetryCount;
     private TopicContainerMap topicContainerMap = TopicContainerMap.empty();
 
     public CosmosDBConfig(ConfigDef config, Map<String, String> parsedConfig) {
         super(config, parsedConfig);
 
-        connEndpoint = this.getString(COSMOS_CONN_ENDPOINT_CONF);
-        connKey = this.getPassword(COSMOS_CONN_KEY_CONF).value();
-        databaseName = this.getString(COSMOS_DATABASE_NAME_CONF);
-        topicContainerMap = TopicContainerMap.deserialize(this.getString(COSMOS_CONTAINER_TOPIC_MAP_CONF));
-        providerName = this.getString(COSMOS_PROVIDER_NAME_CONF);
+        this.connEndpoint = this.getString(COSMOS_CONN_ENDPOINT_CONF);
+        this.connKey = this.getPassword(COSMOS_CONN_KEY_CONF).value();
+        this.databaseName = this.getString(COSMOS_DATABASE_NAME_CONF);
+        this.topicContainerMap = TopicContainerMap.deserialize(this.getString(COSMOS_CONTAINER_TOPIC_MAP_CONF));
+        this.providerName = this.getString(COSMOS_PROVIDER_NAME_CONF);
+        this.bulkModeEnabled = this.getBoolean(COSMOS_SINK_BULK_ENABLED_CONF);
+        this.maxRetryCount = this.getInt(COSMOS_SINK_MAX_RETRY_COUNT);
+        this.gatewayModeEnabled = this.getBoolean(COSMOS_GATEWAY_MODE_ENABLED);
+        this.connectionSharingEnabled = this.getBoolean(COSMOS_CONNECTION_SHARING_ENABLED);
     }
 
     public CosmosDBConfig(Map<String, String> parsedConfig) {
@@ -125,6 +156,33 @@ public class CosmosDBConfig extends AbstractConfig {
                         "none",
                         Importance.MEDIUM,
                         TOLERANCE_ON_ERROR_DOC
+                ).define(
+                        COSMOS_SINK_BULK_ENABLED_CONF,
+                        Type.BOOLEAN,
+                        DEFAULT_COSMOS_SINK_BULK_ENABLED,
+                        Importance.LOW,
+                        COSMOS_SINK_BULK_ENABLED_DOC
+                )
+                .define(
+                        COSMOS_SINK_MAX_RETRY_COUNT,
+                        Type.INT,
+                        DEFAULT_COSMOS_SINK_MAX_RETRY_COUNT,
+                        Importance.HIGH,
+                        COSMOS_SINK_MAX_RETRY_COUNT_DOC
+                )
+                .define(
+                        COSMOS_GATEWAY_MODE_ENABLED,
+                        Type.BOOLEAN,
+                        DEFAULT_COSMOS_GATEWAY_MODE_ENABLED,
+                        Importance.LOW,
+                        COSMOS_GATEWAY_MODE_ENABLED_DOC
+                )
+                .define(
+                        COSMOS_CONNECTION_SHARING_ENABLED,
+                        Type.BOOLEAN,
+                        DEFAULT_COSMOS_CONNECTION_SHARING_ENABLED,
+                        Importance.LOW,
+                        COSMOS_CONNECTION_SHARING_ENABLED_DOC
                 )
                 .defineInternal(
                         COSMOS_PROVIDER_NAME_CONF,
@@ -188,6 +246,22 @@ public class CosmosDBConfig extends AbstractConfig {
         return this.providerName;
     }
 
+    public boolean isBulkModeEnabled() {
+        return this.bulkModeEnabled;
+    }
+
+    public int getMaxRetryCount() {
+        return this.maxRetryCount;
+    }
+
+    public boolean isGatewayModeEnabled() {
+        return gatewayModeEnabled;
+    }
+
+    public boolean isConnectionSharingEnabled() {
+        return this.connectionSharingEnabled;
+    }
+
     public static void validateConnection(Map<String, String> connectorConfigs, Map<String, ConfigValue> configValues) {
         String endpoint = connectorConfigs.get(CosmosDBSinkConfig.COSMOS_CONN_ENDPOINT_CONF);
         String key = connectorConfigs.get(CosmosDBSinkConfig.COSMOS_CONN_KEY_CONF);
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/TopicContainerMap.java b/src/main/java/com/azure/cosmos/kafka/connect/TopicContainerMap.java
index f18eec6..b91ee34 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/TopicContainerMap.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/TopicContainerMap.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect;
 
 import org.apache.commons.collections4.BidiMap;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/BulkWriter.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/BulkWriter.java
new file mode 100644
index 0000000..345ffe6
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/BulkWriter.java
@@ -0,0 +1,150 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.BridgeInternal;
+import com.azure.cosmos.CosmosContainer;
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.ImplementationBridgeHelpers;
+import com.azure.cosmos.implementation.routing.PartitionKeyInternal;
+import com.azure.cosmos.models.CosmosBulkItemResponse;
+import com.azure.cosmos.models.CosmosBulkOperationResponse;
+import com.azure.cosmos.models.CosmosBulkOperations;
+import com.azure.cosmos.models.CosmosItemOperation;
+import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.PartitionKeyDefinition;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument;
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+public class BulkWriter extends SinkWriterBase {
+    private static final Logger logger = LoggerFactory.getLogger(BulkWriter.class);
+
+    private final CosmosContainer cosmosContainer;
+    private final PartitionKeyDefinition partitionKeyDefinition;
+
+    public BulkWriter(CosmosContainer container, int maxRetryCount) {
+        super(maxRetryCount);
+
+        checkNotNull(container, "Argument 'container' can not be null");
+        this.cosmosContainer = container;
+        this.partitionKeyDefinition = container.read().getProperties().getPartitionKeyDefinition();
+    }
+
+    /***
+     * Bulk write the sink records.
+     *
+     * @param sinkRecords the sink records needs to be written.
+     * @return the list of sink write failed operations.
+     */
+    protected SinkWriteResponse writeCore(List<SinkRecord> sinkRecords) {
+
+        SinkWriteResponse sinkWriteResponse = new SinkWriteResponse();
+        if (sinkRecords == null || sinkRecords.isEmpty()) {
+            return sinkWriteResponse;
+        }
+
+        List<CosmosItemOperation> itemOperations = new ArrayList<>();
+        for (SinkRecord sinkRecord : sinkRecords) {
+            CosmosItemOperation cosmosItemOperation = CosmosBulkOperations.getUpsertItemOperation(
+                    sinkRecord.value(),
+                    this.getPartitionKeyValue(sinkRecord.value()),
+                    new SinkOperationContext(sinkRecord));
+
+            itemOperations.add(cosmosItemOperation);
+        }
+
+        Iterable<CosmosBulkOperationResponse<Object>> responseList = cosmosContainer.executeBulkOperations(itemOperations);
+
+        // Non-transient exceptions will be put in the front of the list
+        for (CosmosBulkOperationResponse<Object> bulkOperationResponse : responseList) {
+            SinkOperationContext context = bulkOperationResponse.getOperation().getContext();
+            checkNotNull(context, "sinkOperationContext should not be null");
+
+            SinkRecord sinkRecord = context.getSinkRecord();
+
+            if (bulkOperationResponse.getException() != null
+                    || bulkOperationResponse.getResponse() == null
+                    || !bulkOperationResponse.getResponse().isSuccessStatusCode()) {
+
+                BulkOperationFailedException exception = handleErrorStatusCode(
+                        bulkOperationResponse.getResponse(),
+                        bulkOperationResponse.getException(),
+                        context);
+
+                // Generally we would want to retry for the transient exceptions, and fail-fast for non-transient exceptions
+                // Putting the non-transient exceptions at the front of the list
+                // so later when deciding retry behavior, only examining the first exception will be enough
+                if (ExceptionsHelper.isTransientFailure(exception)) {
+                    sinkWriteResponse.getFailedRecordResponses().add(new SinkOperationFailedResponse(sinkRecord, exception));
+                } else {
+                    sinkWriteResponse.getFailedRecordResponses().add(0, new SinkOperationFailedResponse(sinkRecord, exception));
+                }
+            } else {
+                sinkWriteResponse.getSucceededRecords().add(sinkRecord);
+            }
+        }
+
+        return sinkWriteResponse;
+    }
+
+    private PartitionKey getPartitionKeyValue(Object recordValue) {
+        checkArgument(recordValue instanceof Map, "Argument 'recordValue' is not valid map format.");
+
+        //TODO: examine the code here for sub-partition
+        String partitionKeyPath = StringUtils.join(this.partitionKeyDefinition.getPaths(), "");
+        Map<String, Object> recordMap = (Map<String, Object>) recordValue;
+        Object partitionKeyValue = recordMap.get(partitionKeyPath.substring(1));
+        PartitionKeyInternal partitionKeyInternal = PartitionKeyInternal.fromObjectArray(Collections.singletonList(partitionKeyValue), false);
+
+        return ImplementationBridgeHelpers
+                .PartitionKeyHelper
+                .getPartitionKeyAccessor()
+                .toPartitionKey(partitionKeyInternal);
+    }
+
+    BulkOperationFailedException handleErrorStatusCode(
+            CosmosBulkItemResponse itemResponse,
+            Exception exception,
+            SinkOperationContext sinkOperationContext) {
+
+        int effectiveStatusCode =
+                itemResponse != null
+                        ? itemResponse.getStatusCode()
+                        : (exception != null && exception instanceof CosmosException ? ((CosmosException)exception).getStatusCode() : HttpConstants.StatusCodes.REQUEST_TIMEOUT);
+        int effectiveSubStatusCode =
+                itemResponse != null
+                        ? itemResponse.getSubStatusCode()
+                        : (exception != null && exception instanceof CosmosException ? ((CosmosException)exception).getSubStatusCode() : 0);
+
+        String errorMessage =
+                String.format(
+                        "Request failed with effectiveStatusCode: {%s}, effectiveSubStatusCode: {%s}, kafkaOffset: {%s}, kafkaPartition: {%s}, topic: {%s}",
+                        effectiveStatusCode,
+                        effectiveSubStatusCode,
+                        sinkOperationContext.getKafkaOffset(),
+                        sinkOperationContext.getKafkaPartition(),
+                        sinkOperationContext.getTopic());
+
+
+        return new BulkOperationFailedException(effectiveStatusCode, effectiveSubStatusCode, errorMessage, exception);
+    }
+
+    private static class BulkOperationFailedException extends CosmosException {
+        protected BulkOperationFailedException(int statusCode, int subStatusCode, String message, Throwable cause) {
+            super(statusCode, message, null, cause);
+            BridgeInternal.setSubStatusCode(this, subStatusCode);
+        }
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfig.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfig.java
index 65e81a5..638b477 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfig.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfig.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import com.azure.cosmos.kafka.connect.sink.id.strategy.AbstractIdStrategyConfig;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnector.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnector.java
index e0ca9ef..66a2829 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnector.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnector.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import static com.azure.cosmos.kafka.connect.CosmosDBConfig.validateConnection;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTask.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTask.java
index 0e675df..0e391c6 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTask.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTask.java
@@ -1,30 +1,31 @@
-package com.azure.cosmos.kafka.connect.sink;
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
 
-import java.time.Duration;
-import java.util.Collection;
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
+package com.azure.cosmos.kafka.connect.sink;
 
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosClientBuilder;
 import com.azure.cosmos.CosmosContainer;
-import com.azure.cosmos.CosmosException;
 import com.azure.cosmos.ThrottlingRetryOptions;
-import com.azure.cosmos.implementation.BadRequestException;
 import com.azure.cosmos.kafka.connect.CosmosDBConfig;
 import com.azure.cosmos.kafka.connect.sink.id.strategy.AbstractIdStrategyConfig;
 import com.azure.cosmos.kafka.connect.sink.id.strategy.IdStrategy;
-
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.kafka.connect.data.Struct;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.slf4j.MDC;
 
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
+
 import static com.azure.cosmos.kafka.connect.CosmosDBConfig.TOLERANCE_ON_ERROR_CONFIG;
 
 /**
@@ -34,6 +35,7 @@ public class CosmosDBSinkTask extends SinkTask {
     private static final Logger logger = LoggerFactory.getLogger(CosmosDBSinkTask.class);
     private CosmosClient client = null;
     private CosmosDBSinkConfig config;
+    private final ConcurrentHashMap<String, SinkWriterBase> containerWriterMap = new ConcurrentHashMap<>();
 
     @Override
     public String version() {
@@ -45,15 +47,21 @@ public class CosmosDBSinkTask extends SinkTask {
         logger.trace("Sink task started.");
         this.config = new CosmosDBSinkConfig(map);
 
-        this.client = new CosmosClientBuilder()
-                .endpoint(config.getConnEndpoint())
-                .key(config.getConnKey())
-                .userAgentSuffix(CosmosDBConfig.COSMOS_CLIENT_USER_AGENT_SUFFIX + version())
-                .throttlingRetryOptions(
-                        new ThrottlingRetryOptions()
-                            .setMaxRetryAttemptsOnThrottledRequests(Integer.MAX_VALUE)
-                            .setMaxRetryWaitTime(Duration.ofSeconds((Integer.MAX_VALUE / 1000) - 1)))
-                .buildClient();
+        CosmosClientBuilder cosmosClientBuilder =
+                new CosmosClientBuilder()
+                        .endpoint(config.getConnEndpoint())
+                        .key(config.getConnKey())
+                        .connectionSharingAcrossClientsEnabled(this.config.isConnectionSharingEnabled())
+                        .userAgentSuffix(CosmosDBConfig.COSMOS_CLIENT_USER_AGENT_SUFFIX + version())
+                        .throttlingRetryOptions(
+                                new ThrottlingRetryOptions()
+                                        .setMaxRetryAttemptsOnThrottledRequests(Integer.MAX_VALUE)
+                                        .setMaxRetryWaitTime(Duration.ofSeconds((Integer.MAX_VALUE / 1000) - 1)));
+        if (this.config.isGatewayModeEnabled()) {
+            cosmosClientBuilder.gatewayMode();
+        }
+
+        this.client = cosmosClientBuilder.buildClient();
 
         client.createDatabaseIfNotExists(config.getDatabaseName());
     }
@@ -77,13 +85,28 @@ public class CosmosDBSinkTask extends SinkTask {
         for (Map.Entry<String, List<SinkRecord>> entry : recordsByContainer.entrySet()) {
             String containerName = entry.getKey();
             CosmosContainer container = client.getDatabase(config.getDatabaseName()).getContainer(containerName);
+
+            // get the writer for this container
+            IWriter cosmosdbWriter = this.containerWriterMap.compute(containerName, (name, writer) -> {
+                if (writer == null) {
+                    if (this.config.isBulkModeEnabled()) {
+                        writer = new BulkWriter(container, this.config.getMaxRetryCount());
+                    } else {
+                        writer = new PointWriter(container, this.config.getMaxRetryCount());
+                    }
+                }
+
+                return writer;
+            });
+
+            List<SinkRecord> toBeWrittenRecordList = new ArrayList<>();
             for (SinkRecord record : entry.getValue()) {
                 if (record.key() != null) {
                     MDC.put(String.format("CosmosDbSink-%s", containerName), record.key().toString());
                 }
-                logger.debug("Writing record, value type: {}", record.value().getClass().getName());
-                logger.debug("Key Schema: {}", record.keySchema());
-                logger.debug("Value schema: {}", record.valueSchema());
+                logger.trace("Writing record, value type: {}", record.value().getClass().getName());
+                logger.trace("Key Schema: {}", record.keySchema());
+                logger.trace("Value schema: {}", record.valueSchema());
 
                 Object recordValue;
                 if (record.value() instanceof Struct) {
@@ -93,36 +116,59 @@ public class CosmosDBSinkTask extends SinkTask {
                 }
 
                 maybeInsertId(recordValue, record);
+                toBeWrittenRecordList.add(record);
+            }
+
+            try {
+                SinkWriteResponse response = cosmosdbWriter.write(toBeWrittenRecordList);
+
+                if (response.getFailedRecordResponses().size() == 0) {
+                    logger.debug("Sink write completed, {} records succeeded.", response.getSucceededRecords().size());
 
-                try {
-                    addItemToContainer(container, recordValue);
-                } catch (CosmosException | ConnectException bre) {
-                    sendToDlqIfConfigured(record, bre);
-                } finally {
-                    MDC.clear();
+                } else {
+                    this.sendToDlqIfConfigured(response.getFailedRecordResponses());
                 }
+            }  finally {
+                MDC.clear();
             }
         }
     }
 
+    private boolean isTolerantOnError() {
+        return config.getString(TOLERANCE_ON_ERROR_CONFIG).equalsIgnoreCase("all");
+    }
+
     /**
      * Sends data to a dead letter queue
      *
-     * @param record the kafka record that contains error
+     * @param failedResponses the kafka record that failed to write
      */
-    private void sendToDlqIfConfigured(SinkRecord record, RuntimeException exception) {
-
+    private void sendToDlqIfConfigured(List<SinkOperationFailedResponse> failedResponses) {
         if (context != null && context.errantRecordReporter() != null) {
-            context.errantRecordReporter().report(record, exception);
-            if (!config.getString(TOLERANCE_ON_ERROR_CONFIG).equalsIgnoreCase("all")) {
-                throw exception;
+            for (SinkOperationFailedResponse sinkRecordResponse : failedResponses) {
+                context.errantRecordReporter().report(sinkRecordResponse.getSinkRecord(), sinkRecordResponse.getException());
             }
+        }
+
+        StringBuilder errorMessage = new StringBuilder();
+        for (SinkOperationFailedResponse failedResponse : failedResponses) {
+            errorMessage
+                    .append(
+                        String.format(
+                                "Unable to write record to CosmosDB: {%s}, value schema {%s}, exception {%s}",
+                                failedResponse.getSinkRecord().key(),
+                                failedResponse.getSinkRecord().valueSchema(),
+                                failedResponse.getException().toString()))
+                    .append("\n");
+        }
+
+        if (this.isTolerantOnError()) {
+            logger.error(
+                    "Could not upload record to CosmosDb, but tolerance is set to all. Error message: {}",
+                    errorMessage);
+            return;
         } else {
-            if (config.getString(TOLERANCE_ON_ERROR_CONFIG).equalsIgnoreCase("all")) {
-                logger.error("Could not upload record to CosmosDb, but tolerance is set to all.", exception);
-            } else {
-                throw new CosmosDBWriteException(record, exception);
-            }
+            throw new CosmosDBWriteException(errorMessage.toString());
         }
     }
 
@@ -135,20 +181,14 @@ public class CosmosDBSinkTask extends SinkTask {
         recordMap.put(AbstractIdStrategyConfig.ID, idStrategy.generateId(sinkRecord));
     }
 
-    private void addItemToContainer(CosmosContainer container, Object recordValue) {
-        container.upsertItem(recordValue);
-    }
-
     @Override
     public void stop() {
         logger.trace("Stopping CosmosDB sink task");
 
         if (client != null) {
             client.close();
-            client = null;
         }
 
         client = null;
-
     }
 }
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBWriteException.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBWriteException.java
index e14821a..d682619 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBWriteException.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBWriteException.java
@@ -1,7 +1,9 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import org.apache.kafka.connect.errors.ConnectException;
-import org.apache.kafka.connect.sink.SinkRecord;
 
 /**
  * Exception thrown when an attempt to write a message to CosmosDB has failed.
@@ -12,7 +14,7 @@ public class CosmosDBWriteException extends ConnectException {
      */
     private static final long serialVersionUID = 1L;
 
-    public CosmosDBWriteException(SinkRecord record, Throwable cause) {
-        super("Unable to write record to CosmosDB: " + record.key() + " (value schema:" + record.valueSchema(), cause);
+    public CosmosDBWriteException(String message) {
+        super(message);
     }
 }
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/ExceptionsHelper.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/ExceptionsHelper.java
new file mode 100644
index 0000000..da38762
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/ExceptionsHelper.java
@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.implementation.HttpConstants;
+
+public class ExceptionsHelper {
+    public static boolean isTransientFailure(int statusCode, int substatusCode) {
+        return statusCode == HttpConstants.StatusCodes.GONE
+                || statusCode == HttpConstants.StatusCodes.SERVICE_UNAVAILABLE
+                || statusCode == HttpConstants.StatusCodes.INTERNAL_SERVER_ERROR
+                || statusCode == HttpConstants.StatusCodes.REQUEST_TIMEOUT
+                || (statusCode == HttpConstants.StatusCodes.NOTFOUND && substatusCode == HttpConstants.SubStatusCodes.READ_SESSION_NOT_AVAILABLE);
+
+    }
+
+    public static boolean isTransientFailure(Exception e) {
+        if (e instanceof CosmosException) {
+            return isTransientFailure(((CosmosException) e).getStatusCode(), ((CosmosException) e).getSubStatusCode());
+        }
+
+        return false;
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/IWriter.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/IWriter.java
new file mode 100644
index 0000000..04bfe22
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/IWriter.java
@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import java.util.List;
+
+public interface IWriter {
+    SinkWriteResponse write(List<SinkRecord> sinkRecords);
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/PointWriter.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/PointWriter.java
new file mode 100644
index 0000000..4062915
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/PointWriter.java
@@ -0,0 +1,52 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.CosmosContainer;
+import com.azure.cosmos.CosmosException;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import java.util.List;
+
+import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkNotNull;
+
+public class PointWriter extends SinkWriterBase {
+    private final CosmosContainer container;
+
+    public PointWriter(CosmosContainer container, int maxRetryCount) {
+        super(maxRetryCount);
+
+        checkNotNull(container, "Argument 'container' can not be null");
+        this.container = container;
+    }
+
+    @Override
+    protected SinkWriteResponse writeCore(List<SinkRecord> sinkRecords) {
+        checkNotNull(sinkRecords, "Argument 'sinkRecords' should not be null");
+        SinkWriteResponse sinkWriteResponse = new SinkWriteResponse();
+
+        for (SinkRecord sinkRecord : sinkRecords) {
+            try {
+                container.upsertItem(sinkRecord.value());
+                sinkWriteResponse.getSucceededRecords().add(sinkRecord);
+            } catch (CosmosException cosmosException) {
+                // Generally we would want to retry for the transient exceptions, and fail-fast for non-transient exceptions
+                // Putting the non-transient exceptions at the front of the list
+                // so later when deciding retry behavior, only examining the first exception will be enough
+                if (ExceptionsHelper.isTransientFailure(cosmosException)) {
+                    sinkWriteResponse
+                            .getFailedRecordResponses()
+                            .add(new SinkOperationFailedResponse(sinkRecord, cosmosException));
+                } else {
+                    sinkWriteResponse
+                            .getFailedRecordResponses()
+                            .add(0, new SinkOperationFailedResponse(sinkRecord, cosmosException));
+                }
+
+            }
+        }
+
+        return sinkWriteResponse;
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationContext.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationContext.java
new file mode 100644
index 0000000..24c673d
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationContext.java
@@ -0,0 +1,30 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+
+public class SinkOperationContext {
+    private final SinkRecord sinkRecord;
+
+    public SinkOperationContext(SinkRecord sinkRecord) {
+        this.sinkRecord = sinkRecord;
+    }
+
+    public SinkRecord getSinkRecord() {
+        return this.sinkRecord;
+    }
+
+    public long getKafkaOffset() {
+        return this.sinkRecord.kafkaOffset();
+    }
+
+    public Integer getKafkaPartition() {
+        return this.sinkRecord.kafkaPartition();
+    }
+
+    public String getTopic() {
+        return this.sinkRecord.topic();
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationFailedResponse.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationFailedResponse.java
new file mode 100644
index 0000000..36a7964
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkOperationFailedResponse.java
@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.CosmosException;
+import org.apache.kafka.connect.sink.SinkRecord;
+
+public class SinkOperationFailedResponse {
+    private final SinkRecord sinkRecord;
+    private final Exception exception;
+
+    public SinkOperationFailedResponse(SinkRecord sinkRecord, CosmosException cosmosException) {
+        this.sinkRecord = sinkRecord;
+        this.exception = cosmosException;
+    }
+
+    public SinkRecord getSinkRecord() {
+        return this.sinkRecord;
+    }
+
+    public Exception getException() {
+        return this.exception;
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriteResponse.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriteResponse.java
new file mode 100644
index 0000000..8237e34
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriteResponse.java
@@ -0,0 +1,31 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class SinkWriteResponse {
+    private final List<SinkRecord> succeededRecords;
+    private List<SinkOperationFailedResponse> failedRecordResponses;
+
+    public SinkWriteResponse() {
+        succeededRecords = new ArrayList<>();
+        failedRecordResponses = new ArrayList<>();
+    }
+
+    public List<SinkRecord> getSucceededRecords() {
+        return succeededRecords;
+    }
+
+    public List<SinkOperationFailedResponse> getFailedRecordResponses() {
+        return failedRecordResponses;
+    }
+
+    public void setFailedRecordResponses(List<SinkOperationFailedResponse> failedRecordResponses) {
+        this.failedRecordResponses = failedRecordResponses;
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriterBase.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriterBase.java
new file mode 100644
index 0000000..0d8a4d4
--- /dev/null
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/SinkWriterBase.java
@@ -0,0 +1,61 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+public abstract class SinkWriterBase implements IWriter {
+    private static final Logger logger = LoggerFactory.getLogger(SinkWriterBase.class);
+    private final int maxRetryCount;
+
+    public SinkWriterBase(int maxRetryCount) {
+        this.maxRetryCount = maxRetryCount;
+    }
+
+    protected abstract SinkWriteResponse writeCore(List<SinkRecord> sinkRecords);
+
+    @Override
+    public SinkWriteResponse write(List<SinkRecord> sinkRecords) {
+
+        SinkWriteResponse sinkWriteResponse = writeCore(sinkRecords);
+        int retryCount = 0;
+
+        List<SinkRecord> toBeRetriedRecords;
+        while (shouldRetry(retryCount, sinkWriteResponse)) {
+            retryCount++;
+            toBeRetriedRecords = sinkWriteResponse.getFailedRecordResponses().stream().map(SinkOperationFailedResponse::getSinkRecord).collect(Collectors.toList());
+            SinkWriteResponse retryResponse = writeCore(toBeRetriedRecords);
+            sinkWriteResponse.getSucceededRecords().addAll(retryResponse.getSucceededRecords());
+            sinkWriteResponse.setFailedRecordResponses(retryResponse.getFailedRecordResponses());
+        }
+
+        return sinkWriteResponse;
+    }
+
+    private boolean shouldRetry(int currentRetryCount, SinkWriteResponse response) {
+        if (response == null || response.getFailedRecordResponses().size() == 0) {
+            // there is no failed operation
+            return false;
+        }
+
+        if (currentRetryCount >= this.maxRetryCount) {
+            logger.warn("Exhausted all the retries, will not retry anymore.");
+            return false;
+        }
+
+        // If there are any non-transient exceptions, then retry will NOT happen
+        // for optimization purpose, the non-transient exception will be put in the front
+        // so it will be enough to only examine the first record in the list
+        if (!ExceptionsHelper.isTransientFailure(response.getFailedRecordResponses().get(0).getException())) {
+            return false;
+        }
+
+        return true;
+    }
+}
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMap.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMap.java
index e54cd3f..2b515b7 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMap.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMap.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import java.util.ArrayList;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfig.java b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfig.java
index 09aa5bb..5af3613 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfig.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfig.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import java.util.Arrays;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnector.java b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnector.java
index 440de0d..21f1925 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnector.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnector.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import static com.azure.cosmos.kafka.connect.CosmosDBConfig.validateConnection;
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTask.java b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTask.java
index 0c2e393..6eaa734 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTask.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTask.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import static java.lang.Thread.sleep;
@@ -55,7 +58,7 @@ public class CosmosDBSourceTask extends SourceTask {
         this.queue = new LinkedTransferQueue<>();
 
         logger.info("Creating the client.");
-        client = getCosmosClient();
+        client = getCosmosClient(config);
 
         // Initialize the database, feed and lease containers
         CosmosAsyncDatabase database = client.getDatabase(config.getDatabaseName());
@@ -238,16 +241,22 @@ public class CosmosDBSourceTask extends SourceTask {
         }
     }
 
-    private CosmosAsyncClient getCosmosClient() {
+    private CosmosAsyncClient getCosmosClient(CosmosDBSourceConfig config) {
         logger.info("Creating Cosmos Client.");
 
-        return new CosmosClientBuilder()
+        CosmosClientBuilder cosmosClientBuilder = new CosmosClientBuilder()
                 .endpoint(config.getConnEndpoint())
                 .key(config.getConnKey())
                 .consistencyLevel(ConsistencyLevel.SESSION)
                 .contentResponseOnWriteEnabled(true)
-                .userAgentSuffix(CosmosDBConfig.COSMOS_CLIENT_USER_AGENT_SUFFIX + version())
-                .buildAsyncClient();
+                .connectionSharingAcrossClientsEnabled(config.isConnectionSharingEnabled())
+                .userAgentSuffix(CosmosDBConfig.COSMOS_CLIENT_USER_AGENT_SUFFIX + version());
+
+        if (config.isGatewayModeEnabled()) {
+            cosmosClientBuilder.gatewayMode();
+        }
+
+        return cosmosClientBuilder.buildAsyncClient();
     }
 
     private ChangeFeedProcessor getChangeFeedProcessor(String hostName, CosmosAsyncContainer feedContainer, CosmosAsyncContainer leaseContainer) {
diff --git a/src/main/java/com/azure/cosmos/kafka/connect/source/JsonToStruct.java b/src/main/java/com/azure/cosmos/kafka/connect/source/JsonToStruct.java
index ff30f60..622c86b 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/source/JsonToStruct.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/source/JsonToStruct.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import static java.lang.String.format;


Test patch:
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/CosmosConfigTest.java b/src/test/java/com/azure/cosmos/kafka/connect/CosmosConfigTest.java
index 523fefa..7008d4c 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/CosmosConfigTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/CosmosConfigTest.java
@@ -1,13 +1,14 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect;
 
-import org.apache.kafka.common.config.ConfigException;
 import org.junit.Test;
 
-import java.net.URISyntaxException;
-import java.net.UnknownHostException;
 import java.util.HashMap;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 
 public class CosmosConfigTest {
     private static final String COSMOS_URL = "https://<cosmosinstance-name>.documents.azure.com:443/";
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/IntegrationTest.java b/src/test/java/com/azure/cosmos/kafka/connect/IntegrationTest.java
index 1393fff..70e7648 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/IntegrationTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/IntegrationTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect;
 
 /**
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/TopicContainerMapTest.java b/src/test/java/com/azure/cosmos/kafka/connect/TopicContainerMapTest.java
index 28b8476..f3f1f26 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/TopicContainerMapTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/TopicContainerMapTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect;
 
 import org.junit.Test;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/BulkWriterTests.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/BulkWriterTests.java
new file mode 100644
index 0000000..ed4a7de
--- /dev/null
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/BulkWriterTests.java
@@ -0,0 +1,213 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.CosmosContainer;
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.implementation.BadRequestException;
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.RequestTimeoutException;
+import com.azure.cosmos.models.CosmosBulkItemResponse;
+import com.azure.cosmos.models.CosmosBulkOperationResponse;
+import com.azure.cosmos.models.CosmosBulkOperations;
+import com.azure.cosmos.models.CosmosContainerProperties;
+import com.azure.cosmos.models.CosmosContainerResponse;
+import com.azure.cosmos.models.CosmosItemOperation;
+import com.azure.cosmos.models.PartitionKey;
+import com.azure.cosmos.models.PartitionKeyDefinition;
+import com.google.common.collect.Iterators;
+import org.apache.kafka.connect.data.ConnectSchema;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import static junit.framework.TestCase.assertEquals;
+import static junit.framework.TestCase.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+public class BulkWriterTests {
+    private final int MAX_RETRY_COUNT = 2;
+    private final String TOPIC_NAME = "testtopic";
+
+    private CosmosContainer container;
+    private BulkWriter bulkWriter;
+
+    @Before
+    public void setup(){
+        container = Mockito.mock(CosmosContainer.class);
+
+        CosmosContainerResponse mockedContainerResponse = Mockito.mock(CosmosContainerResponse.class);
+        Mockito.when(container.read()).thenReturn(mockedContainerResponse);
+        CosmosContainerProperties mockedContainerProperties = Mockito.mock(CosmosContainerProperties.class);
+        Mockito.when(mockedContainerResponse.getProperties()).thenReturn(mockedContainerProperties);
+        PartitionKeyDefinition mockedPartitionKeyDefinition = Mockito.mock(PartitionKeyDefinition.class);
+        Mockito.when(mockedContainerProperties.getPartitionKeyDefinition()).thenReturn(mockedPartitionKeyDefinition);
+        Mockito.when(mockedPartitionKeyDefinition.getPaths()).thenReturn(Arrays.asList("/id"));
+
+        bulkWriter = new BulkWriter(container, MAX_RETRY_COUNT);
+    }
+
+    @Test
+    public void testBulkWriteSuccess() {
+        String record1Id = UUID.randomUUID().toString();
+        String record2Id = UUID.randomUUID().toString();
+        SinkRecord record1 = createSinkRecord(record1Id);
+        SinkRecord record2 = createSinkRecord(record2Id);
+
+        // setup successful item response
+        List<CosmosBulkOperationResponse<Object>> mockedBulkOperationResponseList = new ArrayList<>();
+        mockedBulkOperationResponseList.add(mockSuccessfulBulkOperationResponse(record1, record1Id));
+        mockedBulkOperationResponseList.add(mockSuccessfulBulkOperationResponse(record2, record2Id));
+
+        Mockito.when(container.executeBulkOperations(any())).thenReturn(() -> mockedBulkOperationResponseList.iterator());
+
+        SinkWriteResponse response = bulkWriter.write(Arrays.asList(record1, record2));
+        assertEquals(2, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(record2, response.getSucceededRecords().get(1));
+        assertEquals(0, response.getFailedRecordResponses().size());
+    }
+
+    @Test
+    public void testBulkWriteWithNonTransientException() {
+        String record1Id = UUID.randomUUID().toString();
+        String record2Id = UUID.randomUUID().toString();
+        SinkRecord record1 = createSinkRecord(record1Id);
+        SinkRecord record2 = createSinkRecord(record2Id);
+
+        List<CosmosBulkOperationResponse<Object>> mockedBulkOperationResponseList = new ArrayList<>();
+        mockedBulkOperationResponseList.add(mockSuccessfulBulkOperationResponse(record1, record1Id));
+        mockedBulkOperationResponseList.add(mockFailedBulkOperationResponse(record2, record2Id, new BadRequestException("Test")));
+
+        Mockito.when(container.executeBulkOperations(any())).thenReturn(() -> mockedBulkOperationResponseList.iterator());
+
+        SinkWriteResponse response = bulkWriter.write(Arrays.asList(record1, record2));
+        // Validate record 1 succeeded
+        assertEquals(1, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(1, response.getFailedRecordResponses().size());
+
+        // Validate record2 failed
+        assertEquals(record2, response.getFailedRecordResponses().get(0).getSinkRecord());
+        assertTrue(response.getFailedRecordResponses().get(0).getException() instanceof CosmosException);
+        assertEquals(HttpConstants.StatusCodes.BADREQUEST, ((CosmosException)response.getFailedRecordResponses().get(0).getException()).getStatusCode());
+
+        ArgumentCaptor<Iterable<CosmosItemOperation>> parameters = ArgumentCaptor.forClass(Iterable.class);
+        verify(container, times(1)).executeBulkOperations(parameters.capture());
+
+        Iterator<CosmosItemOperation> bulkExecutionParameters = parameters.getValue().iterator();
+        assertEquals(2, Iterators.size(bulkExecutionParameters));
+    }
+
+    @Test
+    public void testBulkWriteSucceededWithTransientException() {
+        String record1Id = UUID.randomUUID().toString();
+        String record2Id = UUID.randomUUID().toString();
+        SinkRecord record1 = createSinkRecord(record1Id);
+        SinkRecord record2 = createSinkRecord(record2Id);
+
+        CosmosBulkOperationResponse<Object> successfulResponseForRecord1 = mockSuccessfulBulkOperationResponse(record1, record1Id);
+        CosmosBulkOperationResponse<Object> failedResponseForRecord2 = mockFailedBulkOperationResponse(record2, record2Id, new RequestTimeoutException());
+        CosmosBulkOperationResponse<Object> successfulResponseForRecord2 = mockSuccessfulBulkOperationResponse(record2, record2Id);
+
+        Mockito.when(container.executeBulkOperations(any()))
+                .thenReturn(() -> Arrays.asList(successfulResponseForRecord1, failedResponseForRecord2).iterator())
+                .thenReturn(() -> Arrays.asList(failedResponseForRecord2).iterator())
+                .thenReturn(() -> Arrays.asList(successfulResponseForRecord2).iterator());
+
+        SinkWriteResponse response = bulkWriter.write(Arrays.asList(record1, record2));
+
+        assertEquals(2, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(record2, response.getSucceededRecords().get(1));
+        assertEquals(0, response.getFailedRecordResponses().size());
+
+        ArgumentCaptor<Iterable<CosmosItemOperation>> parameters = ArgumentCaptor.forClass(Iterable.class);
+        verify(container, times(3)).executeBulkOperations(parameters.capture());
+
+        List<Iterable<CosmosItemOperation>> allParameters = parameters.getAllValues();
+        assertEquals(3, allParameters.size());
+        assertEquals(2, Iterators.size(allParameters.get(0).iterator()));
+        assertEquals(1, Iterators.size(allParameters.get(1).iterator()));
+        assertEquals(1, Iterators.size(allParameters.get(2).iterator()));
+    }
+
+
+    @Test
+    public void testBulkWriteFailedWithTransientException() {
+        String record1Id = UUID.randomUUID().toString();
+        String record2Id = UUID.randomUUID().toString();
+        SinkRecord record1 = createSinkRecord(record1Id);
+        SinkRecord record2 = createSinkRecord(record2Id);
+
+        CosmosBulkOperationResponse<Object> successfulResponseForRecord1 = mockSuccessfulBulkOperationResponse(record1, record1Id);
+        CosmosBulkOperationResponse<Object> failedResponseForRecord2 = mockFailedBulkOperationResponse(record2, record2Id, new RequestTimeoutException());
+
+        Mockito.when(container.executeBulkOperations(any()))
+                .thenReturn(() -> Arrays.asList(successfulResponseForRecord1, failedResponseForRecord2).iterator())
+                .thenReturn(() -> Arrays.asList(failedResponseForRecord2).iterator())
+                .thenReturn(() -> Arrays.asList(failedResponseForRecord2).iterator());
+
+        SinkWriteResponse response = bulkWriter.write(Arrays.asList(record1, record2));
+
+        assertEquals(1, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(1, response.getFailedRecordResponses().size());
+        assertEquals(record2, response.getFailedRecordResponses().get(0).getSinkRecord());
+        assertTrue(response.getFailedRecordResponses().get(0).getException() instanceof CosmosException);
+        assertEquals(HttpConstants.StatusCodes.REQUEST_TIMEOUT, ((CosmosException)response.getFailedRecordResponses().get(0).getException()).getStatusCode());
+    }
+
+    private SinkRecord createSinkRecord(String id) {
+        Schema stringSchema = new ConnectSchema(Schema.Type.STRING);
+        Schema mapSchema = new ConnectSchema(Schema.Type.MAP);
+        Map<String, String> map = new HashMap<>();
+        map.put("foo", "baaarrrrrgh");
+        map.put("id", id);
+
+        return new SinkRecord(TOPIC_NAME, 1, stringSchema, "nokey", mapSchema, map, 0L);
+    }
+
+    private CosmosBulkOperationResponse mockSuccessfulBulkOperationResponse(SinkRecord sinkRecord, String partitionKeyValue) {
+        CosmosBulkItemResponse mockedItemResponse = Mockito.mock(CosmosBulkItemResponse.class);
+        Mockito.when(mockedItemResponse.isSuccessStatusCode()).thenReturn(true);
+
+        CosmosItemOperation itemOperation = CosmosBulkOperations.getUpsertItemOperation(
+                sinkRecord,
+                new PartitionKey(partitionKeyValue),
+                new SinkOperationContext(sinkRecord));
+
+        CosmosBulkOperationResponse<Object> mockedBulkOptionResponse = Mockito.mock(CosmosBulkOperationResponse.class);
+        Mockito.when(mockedBulkOptionResponse.getResponse()).thenReturn(mockedItemResponse);
+        Mockito.when(mockedBulkOptionResponse.getOperation()).thenReturn(itemOperation);
+
+        return mockedBulkOptionResponse;
+    }
+
+    private CosmosBulkOperationResponse mockFailedBulkOperationResponse(SinkRecord sinkRecord, String partitionKeyValue, Exception exception) {
+        CosmosItemOperation itemOperation = CosmosBulkOperations.getUpsertItemOperation(
+                sinkRecord,
+                new PartitionKey(partitionKeyValue),
+                new SinkOperationContext(sinkRecord));
+
+        CosmosBulkOperationResponse<Object> mockedBulkOptionResponse = Mockito.mock(CosmosBulkOperationResponse.class);
+        Mockito.when(mockedBulkOptionResponse.getException()).thenReturn(exception);
+        Mockito.when(mockedBulkOptionResponse.getOperation()).thenReturn(itemOperation);
+
+        return mockedBulkOptionResponse;
+    }
+}
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfigTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfigTest.java
index be414a3..c1fa085 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfigTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConfigTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import com.azure.cosmos.kafka.connect.sink.id.strategy.AbstractIdStrategyConfig;
@@ -12,13 +15,20 @@ import static org.junit.Assert.*;
 public class CosmosDBSinkConfigTest {
 
     private static final String COSMOS_URL = "https://<cosmosinstance-name>.documents.azure.com:443/";
+    private static final String COSMOS_KEY = "mykey";
+    private static final String COSMOS_DATABASE_NAME = "mydb";
+    private static final String COSMOS_CONTAINER_NAME = "mycontainer";
+    private static final String TOPIC_NAME = "mytopic";
+    private static final int DEFAULT_MAX_RETRy = 10;
+    private static final boolean DEFAULT_GATEWAY_MODE_ENABLED = false;
+    private static final boolean DEFAULT_CONNECTION_SHARING_ENABLED = false;
 
     public static HashMap<String, String> setupConfigs() {
         HashMap<String, String> configs = new HashMap<>();
         configs.put(CosmosDBSinkConfig.COSMOS_CONN_ENDPOINT_CONF, COSMOS_URL);
-        configs.put(CosmosDBSinkConfig.COSMOS_CONN_KEY_CONF, "mykey");
-        configs.put(CosmosDBSinkConfig.COSMOS_DATABASE_NAME_CONF, "mydb");
-        configs.put(CosmosDBSinkConfig.COSMOS_CONTAINER_TOPIC_MAP_CONF, "mytopic#mycontainer");
+        configs.put(CosmosDBSinkConfig.COSMOS_CONN_KEY_CONF, COSMOS_KEY);
+        configs.put(CosmosDBSinkConfig.COSMOS_DATABASE_NAME_CONF, COSMOS_DATABASE_NAME);
+        configs.put(CosmosDBSinkConfig.COSMOS_CONTAINER_TOPIC_MAP_CONF, TOPIC_NAME + "#" + COSMOS_CONTAINER_NAME);
         configs.put(AbstractIdStrategyConfig.ID_STRATEGY, KafkaMetadataStrategy.class.getName());
         return configs;
     }
@@ -29,9 +39,9 @@ public class CosmosDBSinkConfigTest {
         CosmosDBSinkConfig config = new CosmosDBSinkConfig(setupConfigs());
         assertNotNull(config);
         assertEquals(COSMOS_URL, config.getConnEndpoint());
-        assertEquals("mykey", config.getConnKey());
-        assertEquals("mydb", config.getDatabaseName());
-        assertEquals("mycontainer", config.getTopicContainerMap().getContainerForTopic("mytopic").get());
+        assertEquals(COSMOS_KEY, config.getConnKey());
+        assertEquals(COSMOS_DATABASE_NAME, config.getDatabaseName());
+        assertEquals(COSMOS_CONTAINER_NAME, config.getTopicContainerMap().getContainerForTopic(TOPIC_NAME).get());
     }
 
     @Test
@@ -55,4 +65,60 @@ public class CosmosDBSinkConfigTest {
             new CosmosDBSinkConfig(settings);
         });
     }
+
+    @Test
+    public void bulkModeTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate by default bulk mode is enabled
+        CosmosDBSinkConfig config = new CosmosDBSinkConfig(settings);
+        assertTrue(config.isBulkModeEnabled());
+
+        // validate bulk mode is false
+        settings.put(CosmosDBSinkConfig.COSMOS_SINK_BULK_ENABLED_CONF, "false");
+        config = new CosmosDBSinkConfig(settings);
+        assertFalse(config.isBulkModeEnabled());
+    }
+
+    @Test
+    public void maxRetryCountTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate default max retry count
+        CosmosDBSinkConfig config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.getMaxRetryCount(), DEFAULT_MAX_RETRy);
+
+        // validate configured max retry count
+        settings.put(CosmosDBSinkConfig.COSMOS_SINK_MAX_RETRY_COUNT, "3");
+        config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.getMaxRetryCount(), 3);
+    }
+
+    @Test
+    public void gatewayModeEnabledTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate default max retry count
+        CosmosDBSinkConfig config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.isGatewayModeEnabled(), DEFAULT_GATEWAY_MODE_ENABLED);
+
+        // validate configured max retry count
+        settings.put(CosmosDBSinkConfig.COSMOS_GATEWAY_MODE_ENABLED, "true");
+        config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.isGatewayModeEnabled(), true);
+    }
+
+    @Test
+    public void connectionSharingEnabledTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate default max retry count
+        CosmosDBSinkConfig config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.isConnectionSharingEnabled(), DEFAULT_CONNECTION_SHARING_ENABLED);
+
+        // validate configured max retry count
+        settings.put(CosmosDBSinkConfig.COSMOS_CONNECTION_SHARING_ENABLED, "true");
+        config = new CosmosDBSinkConfig(settings);
+        assertEquals(config.isConnectionSharingEnabled(), true);
+    }
 }
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorConfigTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorConfigTest.java
index 463d8d2..8aa450a 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorConfigTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorConfigTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import org.apache.kafka.common.config.ConfigDef;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorTest.java
index ed686d9..9bb1cef 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkConnectorTest.java
@@ -1,22 +1,22 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
-import com.azure.cosmos.kafka.connect.CosmosDBConfig;
 import com.azure.cosmos.kafka.connect.CosmosDBConfig.CosmosClientBuilder;
 import com.google.common.collect.ImmutableMap;
 import org.apache.kafka.common.config.Config;
-import org.apache.kafka.common.config.ConfigException;
 import org.apache.kafka.common.config.ConfigValue;
 import org.junit.Test;
+import org.mockito.MockedStatic;
 
 import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
-import org.mockito.MockedStatic;
 
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.collection.IsEmptyCollection.empty;
-import static org.junit.Assert.assertThrows;
 import static org.mockito.AdditionalAnswers.answerVoid;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mockStatic;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTest.java
index a6d1293..ccfd54d 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTest.java
@@ -1,14 +1,12 @@
-package com.azure.cosmos.kafka.connect.sink;
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
 
+package com.azure.cosmos.kafka.connect.sink;
 
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosContainer;
 import com.azure.cosmos.CosmosDatabase;
 import com.azure.cosmos.implementation.BadRequestException;
-import com.fasterxml.jackson.core.JsonParseException;
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.commons.lang3.reflect.FieldUtils;
 import org.apache.kafka.connect.data.ConnectSchema;
 import org.apache.kafka.connect.data.Schema;
@@ -16,24 +14,44 @@ import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.mockito.MockedConstruction;
 import org.mockito.Mockito;
 
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
+import static junit.framework.TestCase.assertEquals;
 import static junit.framework.TestCase.assertNotNull;
 import static junit.framework.TestCase.fail;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.mockConstructionWithAnswer;
+import static org.mockito.Mockito.when;
 
+@RunWith(Parameterized.class)
 public class CosmosDBSinkTaskTest {
     private final String topicName = "testtopic";
     private final String containerName = "container666";
     private final String databaseName = "fakeDatabase312";
+    private final boolean isBulkModeEnabled;
+
     private CosmosDBSinkTask testTask;
     private CosmosClient mockCosmosClient;
     private CosmosContainer mockContainer;
 
+    @Parameterized.Parameters
+    public static Iterable<? extends Object> bulkModeEnabledProvider() {
+        return Arrays.asList(true, false);
+    }
+
+    public CosmosDBSinkTaskTest(boolean bulkModeEnabled) {
+        this.isBulkModeEnabled = bulkModeEnabled;
+    }
+
     @Before
     public void setup() throws IllegalAccessException {
         testTask = new CosmosDBSinkTask();
@@ -42,6 +60,7 @@ public class CosmosDBSinkTaskTest {
         Map<String, String> settingAssignment = CosmosDBSinkConfigTest.setupConfigs();
         settingAssignment.put(CosmosDBSinkConfig.COSMOS_CONTAINER_TOPIC_MAP_CONF, topicName + "#" + containerName);
         settingAssignment.put(CosmosDBSinkConfig.COSMOS_DATABASE_NAME_CONF, databaseName);
+        settingAssignment.put(CosmosDBSinkConfig.COSMOS_SINK_BULK_ENABLED_CONF, String.valueOf(this.isBulkModeEnabled));
         CosmosDBSinkConfig config = new CosmosDBSinkConfig(settingAssignment);
         FieldUtils.writeField(testTask, "config", config, true);
 
@@ -56,39 +75,57 @@ public class CosmosDBSinkTaskTest {
     }
 
     @Test
-    public void testPutPlainTextString() {
+    public void sinkWriteFailed() {
         Schema stringSchema = new ConnectSchema(Schema.Type.STRING);
 
         SinkRecord record = new SinkRecord(topicName, 1, stringSchema, "nokey", stringSchema, "foo", 0L);
         assertNotNull(record.value());
 
-        //Make mock connector to serialize a non-JSON payload
-        when(mockContainer.upsertItem(any())).then((invocation) -> {
-            Object item = invocation.getArgument(0);
-            //Will throw exception:
-            try {
-                JsonNode jsonNode = new ObjectMapper().readTree(item.toString());
-                assertNotNull(jsonNode);
-                return null;
-            } catch (JsonParseException jpe) {
-                throw new BadRequestException("Unable to serialize JSON request", jpe);
-            }
-        });
+        SinkWriteResponse sinkWriteResponse = new SinkWriteResponse();
+        sinkWriteResponse
+                .getFailedRecordResponses()
+                .add(new SinkOperationFailedResponse(record, new BadRequestException("Unable to serialize JSON request")));
+
+        MockedConstruction<? extends SinkWriterBase> mockedWriterConstruction = null;
 
         try {
-            testTask.put(Arrays.asList(record));
-            fail("Expected ConnectException on bad message");
-        } catch (ConnectException ce) {
+            if (this.isBulkModeEnabled) {
+                mockedWriterConstruction = mockConstructionWithAnswer(BulkWriter.class, invocation -> {
+                    if (invocation.getMethod().equals(BulkWriter.class.getMethod("write", List.class))) {
+                        return sinkWriteResponse;
+                    }
+
+                    throw new IllegalStateException("Not implemented for method " + invocation.getMethod().getName());
+                });
+            } else {
+                mockedWriterConstruction = mockConstructionWithAnswer(PointWriter.class, invocation -> {
+                    if (invocation.getMethod().equals(PointWriter.class.getMethod("write", List.class))) {
+                        return sinkWriteResponse;
+                    }
+
+                    throw new IllegalStateException("Not implemented for method " + invocation.getMethod().getName());
+                });
+            }
 
-        } catch (Throwable t) {
-            fail("Expected ConnectException, but got: " + t.getClass().getName());
-        }
+            try {
+                testTask.put(Arrays.asList(record));
+                fail("Expected ConnectException on bad message");
+            } catch (ConnectException ce) {
 
-        verify(mockContainer, times(1)).upsertItem("foo");
+            } catch (Throwable t) {
+                fail("Expected ConnectException, but got: " + t.getClass().getName());
+            }
+
+            assertEquals(1, mockedWriterConstruction.constructed().size());
+        } finally {
+            if (mockedWriterConstruction != null) {
+                mockedWriterConstruction.close();
+            }
+        }
     }
 
     @Test
-    public void testPutMap() {
+    public void sinkWriteSucceeded() {
         Schema stringSchema = new ConnectSchema(Schema.Type.STRING);
         Schema mapSchema = new ConnectSchema(Schema.Type.MAP);
         Map<String, String> map = new HashMap<>();
@@ -96,8 +133,44 @@ public class CosmosDBSinkTaskTest {
 
         SinkRecord record = new SinkRecord(topicName, 1, stringSchema, "nokey", mapSchema, map, 0L);
         assertNotNull(record.value());
-        testTask.put(Arrays.asList(record));
-        verify(mockContainer, times(1)).upsertItem(map);
+
+        SinkWriteResponse sinkWriteResponse = new SinkWriteResponse();
+        sinkWriteResponse.getSucceededRecords().add(record);
+
+        MockedConstruction<? extends SinkWriterBase> mockedWriterConstruction = null;
+        try {
+            if (this.isBulkModeEnabled) {
+                mockedWriterConstruction = mockConstructionWithAnswer(BulkWriter.class, invocation -> {
+                    if (invocation.getMethod().equals(BulkWriter.class.getMethod("write", List.class))) {
+                        return sinkWriteResponse;
+                    }
+
+                    throw new IllegalStateException("Not implemented for method " + invocation.getMethod().getName());
+                });
+            } else {
+                mockedWriterConstruction = mockConstructionWithAnswer(PointWriter.class, invocation -> {
+                    if (invocation.getMethod().equals(PointWriter.class.getMethod("write", List.class))) {
+                        return sinkWriteResponse;
+                    }
+
+                    throw new IllegalStateException("Not implemented for method " + invocation.getMethod().getName());
+                });
+            }
+
+            try {
+                testTask.put(Arrays.asList(record));
+            } catch (ConnectException ce) {
+                fail("Expected sink write succeeded. but got: " + ce.getMessage());
+            } catch (Throwable t) {
+                fail("Expected sink write succeeded, but got: " + t.getClass().getName());
+            }
+
+            assertEquals(1, mockedWriterConstruction.constructed().size());
+        } finally {
+            if (mockedWriterConstruction != null) {
+                mockedWriterConstruction.close();
+            }
+        }
     }
 }
 
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTestNotFails.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTestNotFails.java
index 70286b3..2f83a61 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTestNotFails.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTaskTestNotFails.java
@@ -1,18 +1,16 @@
-package com.azure.cosmos.kafka.connect.sink;
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
 
+package com.azure.cosmos.kafka.connect.sink;
 
 import com.azure.cosmos.CosmosClient;
 import com.azure.cosmos.CosmosContainer;
 import com.azure.cosmos.CosmosDatabase;
 import com.azure.cosmos.implementation.BadRequestException;
-import com.fasterxml.jackson.core.JsonParseException;
 import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.commons.lang3.reflect.FieldUtils;
 import org.apache.kafka.connect.data.ConnectSchema;
 import org.apache.kafka.connect.data.Schema;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.ErrantRecordReporter;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTaskContext;
@@ -21,14 +19,14 @@ import org.junit.Before;
 import org.junit.Test;
 import org.mockito.Mockito;
 
-import java.util.Arrays;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
-import static junit.framework.TestCase.assertNotNull;
-import static junit.framework.TestCase.fail;
-import static org.mockito.Mockito.*;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyString;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class CosmosDBSinkTaskTestNotFails {
     private final String topicName = "testtopic";
@@ -49,6 +47,7 @@ public class CosmosDBSinkTaskTestNotFails {
         settingAssignment.put(CosmosDBSinkConfig.COSMOS_CONTAINER_TOPIC_MAP_CONF, topicName + "#" + containerName);
         settingAssignment.put(CosmosDBSinkConfig.COSMOS_DATABASE_NAME_CONF, databaseName);
         settingAssignment.put(CosmosDBSinkConfig.TOLERANCE_ON_ERROR_CONFIG, "all");
+        settingAssignment.put(CosmosDBSinkConfig.COSMOS_SINK_BULK_ENABLED_CONF, "false");
         CosmosDBSinkConfig config = new CosmosDBSinkConfig(settingAssignment);
         FieldUtils.writeField(testTask, "config", config, true);
 
@@ -61,14 +60,13 @@ public class CosmosDBSinkTaskTestNotFails {
         when(mockContext.errantRecordReporter()).thenReturn(mockErrantReporter);
 
         FieldUtils.writeField(testTask, "client", mockCosmosClient, true);
-
     }
+
     @After()
     public void resetContext() throws IllegalAccessException {
         FieldUtils.writeField(testTask,  "context", null, true);
     }
 
-
     @Test
     public void testPutMapThatFailsDoesNotStopTask() throws JsonProcessingException, IllegalAccessException {
 
@@ -78,6 +76,7 @@ public class CosmosDBSinkTaskTestNotFails {
         SinkRecord record = new SinkRecord(topicName, 1, stringSchema, "nokey", mapSchema, "{", 0L);
         testTask.put(List.of(record));
     }
+
     @Test
     public void testPutMapThatFailsDoesNotStopTaskWithdlq() throws JsonProcessingException, IllegalAccessException {
         FieldUtils.writeField(testTask,  "context", mockContext, true);
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/PointWriterTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/PointWriterTest.java
new file mode 100644
index 0000000..b99f044
--- /dev/null
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/PointWriterTest.java
@@ -0,0 +1,139 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package com.azure.cosmos.kafka.connect.sink;
+
+import com.azure.cosmos.CosmosContainer;
+import com.azure.cosmos.CosmosException;
+import com.azure.cosmos.implementation.BadRequestException;
+import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.RequestTimeoutException;
+import com.azure.cosmos.models.CosmosItemResponse;
+import org.apache.kafka.connect.data.ConnectSchema;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.sink.SinkRecord;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+
+import static junit.framework.Assert.assertTrue;
+import static junit.framework.TestCase.assertEquals;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+public class PointWriterTest {
+    private final int MAX_RETRY_COUNT = 2;
+    private final String TOPIC_NAME = "testtopic";
+
+    private CosmosContainer container;
+    private PointWriter pointWriter;
+
+    @Before
+    public void setup(){
+        container = Mockito.mock(CosmosContainer.class);
+        pointWriter = new PointWriter(container, MAX_RETRY_COUNT);
+    }
+
+    @Test
+    public void testPointWriteSuccess() {
+        SinkRecord record1 = createSinkRecord();
+        SinkRecord record2 = createSinkRecord();
+        CosmosItemResponse<Object> itemResponse = Mockito.mock(CosmosItemResponse.class);
+
+        Mockito.when(container.upsertItem(record1.value())).thenReturn(itemResponse);
+        Mockito.when(container.upsertItem(record2.value())).thenReturn(itemResponse);
+
+        SinkWriteResponse response = pointWriter.write(Arrays.asList(record1, record2));
+        assertEquals(2, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(record2, response.getSucceededRecords().get(1));
+        assertEquals(0, response.getFailedRecordResponses().size());
+    }
+
+    @Test
+    public void testPointWriteWithNonTransientException() {
+        SinkRecord record1 = createSinkRecord();
+        SinkRecord record2 = createSinkRecord();
+
+        CosmosItemResponse<Object> itemResponse = Mockito.mock(CosmosItemResponse.class);
+        Mockito.when(container.upsertItem(record1.value())).thenReturn(itemResponse);
+        Mockito.when(container.upsertItem(record2.value())).thenThrow(new BadRequestException("Test"));
+
+        SinkWriteResponse response = pointWriter.write(Arrays.asList(record1, record2));
+        // Validate record 1 succeeded
+        assertEquals(1, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(1, response.getFailedRecordResponses().size());
+        verify(container, times(1)).upsertItem(record1.value());
+
+        // Validate record2 failed
+        assertEquals(record2, response.getFailedRecordResponses().get(0).getSinkRecord());
+        assertTrue(response.getFailedRecordResponses().get(0).getException() instanceof CosmosException);
+        assertEquals(HttpConstants.StatusCodes.BADREQUEST, ((CosmosException)response.getFailedRecordResponses().get(0).getException()).getStatusCode());
+        verify(container, times(1)).upsertItem(record2.value());
+    }
+
+    @Test
+    public void testPointWriteSucceededWithTransientException() {
+        SinkRecord record1 = createSinkRecord();
+        SinkRecord record2 = createSinkRecord();
+
+        CosmosItemResponse<Object> itemResponse = Mockito.mock(CosmosItemResponse.class);
+        Mockito.when(container.upsertItem(record1.value())).thenReturn(itemResponse);
+        Mockito
+                .when(container.upsertItem(record2.value()))
+                .thenThrow(new RequestTimeoutException())
+                .thenThrow(new RequestTimeoutException())
+                .thenReturn(itemResponse);
+
+        SinkWriteResponse response = pointWriter.write(Arrays.asList(record1, record2));
+
+        assertEquals(2, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(record2, response.getSucceededRecords().get(1));
+        assertEquals(0, response.getFailedRecordResponses().size());
+
+        verify(container, times(1)).upsertItem(record1.value());
+        verify(container, times(3)).upsertItem(record2.value());
+    }
+
+    @Test
+    public void testPointWriteFailedWithTransientException() {
+        SinkRecord record1 = createSinkRecord();
+        SinkRecord record2 = createSinkRecord();
+
+        CosmosItemResponse<Object> itemResponse = Mockito.mock(CosmosItemResponse.class);
+        Mockito.when(container.upsertItem(record1.value())).thenReturn(itemResponse);
+        Mockito
+                .when(container.upsertItem(record2.value()))
+                .thenThrow(new RequestTimeoutException())
+                .thenThrow(new RequestTimeoutException())
+                .thenThrow(new RequestTimeoutException());
+
+        SinkWriteResponse response = pointWriter.write(Arrays.asList(record1, record2));
+
+        assertEquals(1, response.getSucceededRecords().size());
+        assertEquals(record1, response.getSucceededRecords().get(0));
+        assertEquals(record2, response.getFailedRecordResponses().get(0).getSinkRecord());
+        assertTrue(response.getFailedRecordResponses().get(0).getException() instanceof CosmosException);
+        assertEquals(HttpConstants.StatusCodes.REQUEST_TIMEOUT, ((CosmosException)response.getFailedRecordResponses().get(0).getException()).getStatusCode());
+
+        verify(container, times(1)).upsertItem(record1.value());
+        verify(container, times(3)).upsertItem(record2.value());
+    }
+
+    private SinkRecord createSinkRecord() {
+        Schema stringSchema = new ConnectSchema(Schema.Type.STRING);
+        Schema mapSchema = new ConnectSchema(Schema.Type.MAP);
+        Map<String, String> map = new HashMap<>();
+        map.put("foo", "baaarrrrrgh");
+        map.put("id", UUID.randomUUID().toString());
+
+        return new SinkRecord(TOPIC_NAME, 1, stringSchema, "nokey", mapSchema, map, 0L);
+    }
+}
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMapTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMapTest.java
index 86c3771..994800b 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMapTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/StructToJsonMapTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink;
 
 import com.google.common.collect.ImmutableList;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/ProvidedInStrategyTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/ProvidedInStrategyTest.java
index 64c74d9..5ccfacd 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/ProvidedInStrategyTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/ProvidedInStrategyTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink.id.strategy;
 
 import com.google.common.collect.ImmutableList;
@@ -14,7 +17,6 @@ import org.junit.runners.Parameterized;
 import org.mockito.Mock;
 
 import static org.junit.Assert.assertEquals;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.mockito.MockitoAnnotations.initMocks;
 
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/TemplateStrategyTest.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/TemplateStrategyTest.java
index c14dd06..a2b524e 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/TemplateStrategyTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/id/strategy/TemplateStrategyTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink.id.strategy;
 
 import com.google.common.collect.ImmutableMap;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/sink/integration/SinkConnectorIT.java b/src/test/java/com/azure/cosmos/kafka/connect/sink/integration/SinkConnectorIT.java
index 9facf0e..9210c78 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/sink/integration/SinkConnectorIT.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/sink/integration/SinkConnectorIT.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.sink.integration;
 
 import com.azure.cosmos.CosmosClient;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfigTest.java b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfigTest.java
index a2e2dcf..c7141a9 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfigTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfigTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import org.apache.kafka.common.config.ConfigException;
@@ -5,10 +8,15 @@ import org.junit.Test;
 
 import java.util.HashMap;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
 
 public class CosmosDBSourceConfigTest {
     private static final String COSMOS_URL = "https://<cosmosinstance-name>.documents.azure.com:443/";
+    private static final boolean DEFAULT_GATEWAY_MODE_ENABLED = false;
+    private static final boolean DEFAULT_CONNECTION_SHARING_ENABLED = false;
 
     public static HashMap<String, String> setupConfigs() {
         HashMap<String, String> configs = new HashMap<>();
@@ -50,4 +58,32 @@ public class CosmosDBSourceConfigTest {
             new CosmosDBSourceConfig(settings);
         });
     }
+
+    @Test
+    public void gatewayModeEnabledTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate default max retry count
+        CosmosDBSourceConfig config = new CosmosDBSourceConfig(settings);
+        assertEquals(config.isGatewayModeEnabled(), DEFAULT_GATEWAY_MODE_ENABLED);
+
+        // validate configured max retry count
+        settings.put(CosmosDBSourceConfig.COSMOS_GATEWAY_MODE_ENABLED, "true");
+        config = new CosmosDBSourceConfig(settings);
+        assertEquals(config.isGatewayModeEnabled(), true);
+    }
+
+    @Test
+    public void connectionSharingEnabledTest() {
+        HashMap<String, String> settings = setupConfigs();
+
+        // validate default max retry count
+        CosmosDBSourceConfig config = new CosmosDBSourceConfig(settings);
+        assertEquals(config.isConnectionSharingEnabled(), DEFAULT_CONNECTION_SHARING_ENABLED);
+
+        // validate configured max retry count
+        settings.put(CosmosDBSourceConfig.COSMOS_CONNECTION_SHARING_ENABLED, "true");
+        config = new CosmosDBSourceConfig(settings);
+        assertEquals(config.isConnectionSharingEnabled(), true);
+    }
 }
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnectorTest.java b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnectorTest.java
index 669385d..2bf1af8 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnectorTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConnectorTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import org.apache.kafka.common.config.ConfigDef;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTaskTest.java b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTaskTest.java
index 0a38c27..0a481eb 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTaskTest.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceTaskTest.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source;
 
 import static org.mockito.ArgumentMatchers.anyString;
diff --git a/src/test/java/com/azure/cosmos/kafka/connect/source/integration/SourceConnectorIT.java b/src/test/java/com/azure/cosmos/kafka/connect/source/integration/SourceConnectorIT.java
index a349290..e22ebd0 100644
--- a/src/test/java/com/azure/cosmos/kafka/connect/source/integration/SourceConnectorIT.java
+++ b/src/test/java/com/azure/cosmos/kafka/connect/source/integration/SourceConnectorIT.java
@@ -1,3 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
 package com.azure.cosmos.kafka.connect.source.integration;
 
 import com.azure.cosmos.CosmosClient;