Repository: microsoft/kafka-connect-cosmosdb
Pull request #490 (2022-10-05 00:31:06)

enable gateway support
1. Add gateway support
2. Add connectionSharing support for direct mode


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8f5739d..0e74d9b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 #### Key Bug Fixes
 
 #### Other Changes
+* 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
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 2095391..4576ac3 100644
--- a/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java
+++ b/src/main/java/com/azure/cosmos/kafka/connect/CosmosDBConfig.java
@@ -55,6 +55,17 @@ public class CosmosDBConfig extends AbstractConfig {
             "A comma delimited list of Kafka topics mapped to Cosmos containers.\n"
                     + "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;
@@ -80,6 +91,8 @@ public class CosmosDBConfig extends AbstractConfig {
     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();
 
@@ -93,6 +106,8 @@ public class CosmosDBConfig extends AbstractConfig {
         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) {
@@ -155,6 +170,20 @@ public class CosmosDBConfig extends AbstractConfig {
                         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,
                         Type.STRING,
@@ -225,6 +254,14 @@ public class CosmosDBConfig extends AbstractConfig {
         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/sink/CosmosDBSinkTask.java b/src/main/java/com/azure/cosmos/kafka/connect/sink/CosmosDBSinkTask.java
index 1f32dc0..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
@@ -47,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());
     }
@@ -98,9 +104,9 @@ public class CosmosDBSinkTask extends SinkTask {
                 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) {
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 7bc79fd..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
@@ -58,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());
@@ -241,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) {


Test patch:
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 a8c0674..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
@@ -20,6 +20,8 @@ public class CosmosDBSinkConfigTest {
     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<>();
@@ -91,4 +93,32 @@ public class CosmosDBSinkConfigTest {
         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/source/CosmosDBSourceConfigTest.java b/src/test/java/com/azure/cosmos/kafka/connect/source/CosmosDBSourceConfigTest.java
index e33fe08..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
@@ -8,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<>();
@@ -53,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);
+    }
 }