Changed getConnectionString to incorporate SRV fomrat

Author: PushpenderSaini0

docs/MongoDB-batchsink.md

@@ -26,6 +26,9 @@ automatically generated.
 
 **Password:** Password to use to connect to the specified database.
 
+**Connect Using SRV String:** Toggle to determine whether to use an SRV connection string for MongoDB. It can be
+enabled if the MongoDB deployment supports SRV DNS records for connection resolution.
+
 **Connection Arguments:** A list of arbitrary string key/value pairs as connection arguments. See
 [Connection String Options] for a full description of these arguments.
 

docs/MongoDB-batchsource.md

@@ -28,6 +28,9 @@ and use the [MongoDB extended JSON format] to represent non-native JSON data typ
 
 **Password:** Password to use to connect to the specified database.
 
+**Connect Using SRV String:** Toggle to determine whether to use an SRV connection string for MongoDB. It can be
+enabled if the MongoDB deployment supports SRV DNS records for connection resolution.
+
 **Authentication Connection String:** Optional MongoDB connection string to connect to the 'config' database of a
 sharded cluster. It can be omitted if username and password do not differ from the previously provided ones or if
 'config' database does not require authentication.

src/main/java/io/cdap/plugin/MongoDBConfig.java

@@ -50,7 +50,7 @@ public class MongoDBConfig extends PluginConfig {
   @Name(MongoDBConstants.PORT)
   @Description("Port that MongoDB is listening to.")
   @Macro
-  private int port;
+  private Integer port;
 
   @Name(MongoDBConstants.DATABASE)
   @Description("MongoDB database name.")
@@ -76,21 +76,29 @@ public class MongoDBConfig extends PluginConfig {
   @Nullable
   private String password;
 
+  @Name(MongoDBConstants.CONNECT_USING_SRV_STRING)
+  @Description("Toggle to determine whether to use an SRV connection string for MongoDB. It can be " +
+    "enabled if the MongoDB deployment supports SRV DNS records for connection resolution.")
+  @Macro
+  @Nullable
+  private boolean connectUsingSRVString;
+
   @Name(MongoDBConstants.CONNECTION_ARGUMENTS)
   @Description("A list of arbitrary string key/value pairs as connection arguments.")
   @Macro
   @Nullable
   private String connectionArguments;
 
   public MongoDBConfig(String referenceName, String host, int port, String database, String collection, String user,
-                       String password, String connectionArguments) {
+                       String password, boolean connectUsingSRVString, String connectionArguments) {
     this.referenceName = referenceName;
     this.host = host;
     this.port = port;
     this.database = database;
     this.collection = collection;
     this.user = user;
     this.password = password;
+    this.connectUsingSRVString = connectUsingSRVString;
     this.connectionArguments = connectionArguments;
   }
 
@@ -102,7 +110,8 @@ public String getHost() {
     return host;
   }
 
-  public int getPort() {
+  @Nullable
+  public Integer getPort() {
     return port;
   }
 
@@ -124,6 +133,10 @@ public String getPassword() {
     return password;
   }
 
+  public boolean connectUsingSRVString() {
+    return connectUsingSRVString;
+  }
+
   @Nullable
   public String getConnectionArguments() {
     return connectionArguments;
@@ -146,7 +159,8 @@ public void validate() {
     if (!containsMacro(MongoDBConstants.HOST) && Strings.isNullOrEmpty(host)) {
       throw new InvalidConfigPropertyException("Host must be specified", MongoDBConstants.HOST);
     }
-    if (!containsMacro(MongoDBConstants.PORT)) {
+    if ((!containsMacro(MongoDBConstants.CONNECT_USING_SRV_STRING) && !connectUsingSRVString) &&
+      !containsMacro(MongoDBConstants.PORT)) {
       if (port < 1) {
         throw new InvalidConfigPropertyException("Port number must be greater than 0", MongoDBConstants.PORT);
       }
@@ -161,24 +175,32 @@ public void validate() {
 
   /**
    * Constructs a connection string such as: "mongodb://admin:password@localhost:27017/admin.analytics?key=value;"
-   * using host, port, username, password, database, collection and optional connection properties. In the case when
-   * username or password is not provided the connection string will not contain credentials:
+   * using host, port, username, password, database, collection and optional connection properties.
+   * If SRV is enabled, the connection string will use the "mongodb+srv://" protocol instead of "mongodb://".
+   * In the case when username or password is not provided, the connection string will not contain credentials:
    * "mongodb://localhost:27017/admin.analytics?key=value;"
+   * When SRV is not used, the port will be included in the connection string.
    *
    * @return connection string.
    */
   public String getConnectionString() {
-    StringBuilder connectionStringBuilder = new StringBuilder("mongodb://");
+    StringBuilder connectionStringBuilder = new StringBuilder();
+    if (connectUsingSRVString()) {
+      connectionStringBuilder.append("mongodb+srv://");
+    } else {
+      connectionStringBuilder.append("mongodb://");
+    }
     if (!Strings.isNullOrEmpty(user) || !Strings.isNullOrEmpty(password)) {
       connectionStringBuilder.append(user).append(":").append(password).append("@");
     }
-    connectionStringBuilder.append(host).append(":").append(port).append("/")
-      .append(database).append(".").append(collection);
-
+    connectionStringBuilder.append(host);
+    if (!connectUsingSRVString()) {
+      connectionStringBuilder.append(":").append(port);
+    }
+    connectionStringBuilder.append("/").append(database).append(".").append(collection);
     if (!Strings.isNullOrEmpty(connectionArguments)) {
       connectionStringBuilder.append("?").append(connectionArguments);
     }
-
     return connectionStringBuilder.toString();
   }
 

src/main/java/io/cdap/plugin/MongoDBConstants.java

@@ -79,6 +79,11 @@ private MongoDBConstants() {
    */
   public static final String PASSWORD = "password";
 
+  /**
+   * Configuration property name used to specify whether to use an SRV Connection string for MongoDB.
+   */
+  public static final String CONNECT_USING_SRV_STRING = "connectUsingSRVString";
+
   /**
    * Configuration property name used to specify auxiliary MongoDB connection string to authenticate against when
    * constructing splits.

src/main/java/io/cdap/plugin/batch/sink/MongoDBBatchSink.java

@@ -171,8 +171,10 @@ public static class MongoDBSinkConfig extends MongoDBConfig {
     private String idField;
 
     public MongoDBSinkConfig(String referenceName, String host, int port, String database, String collection,
-                             String user, String password, String connectionArguments, String idField) {
-      super(referenceName, host, port, database, collection, user, password, connectionArguments);
+                             String user, String password, boolean connectUsingSRVString,
+                             String connectionArguments, String idField) {
+      super(referenceName, host, port, database, collection, user, password,
+            connectUsingSRVString, connectionArguments);
       this.idField = idField;
     }
 

src/main/java/io/cdap/plugin/batch/source/MongoDBBatchSource.java

@@ -190,9 +190,11 @@ public static class MongoDBSourceConfig extends MongoDBConfig {
     private String authConnectionString;
 
     public MongoDBSourceConfig(String referenceName, String host, int port, String database, String collection,
-                               String user, String password, String connectionArguments, String schema,
+                               String user, String password, boolean connectUsingSRVString,
+                               String connectionArguments, String schema,
                                String inputQuery, String onError, String authConnectionString) {
-      super(referenceName, host, port, database, collection, user, password, connectionArguments);
+      super(referenceName, host, port, database, collection, user, password,
+            connectUsingSRVString, connectionArguments);
       this.schema = schema;
       this.inputQuery = inputQuery;
       this.onError = onError;

src/test/java/io/cdap/plugin/batch/MongoDBConfigBuilder.java

@@ -29,6 +29,7 @@ public class MongoDBConfigBuilder<T extends MongoDBConfigBuilder> {
   protected String collection;
   protected String user;
   protected String password;
+  protected boolean connectUsingSRVString;
   protected String connectionArguments;
 
   public static MongoDBConfigBuilder builder() {
@@ -82,12 +83,18 @@ public T setPassword(String password) {
     return (T) this;
   }
 
+  public T setConnectUsingSRVString(boolean connectUsingSRVString) {
+    this.connectUsingSRVString = connectUsingSRVString;
+    return (T) this;
+  }
+
   public T setConnectionArguments(String connectionArguments) {
     this.connectionArguments = connectionArguments;
     return (T) this;
   }
 
   public MongoDBConfig build() {
-    return new MongoDBConfig(referenceName, host, port, database, collection, user, password, connectionArguments);
+    return new MongoDBConfig(referenceName, host, port, database, collection, user, password,
+                             connectUsingSRVString, connectionArguments);
   }
 }

src/test/java/io/cdap/plugin/batch/MongoDBConfigTest.java

@@ -36,6 +36,7 @@ public class MongoDBConfigTest {
     .setCollection("analytics")
     .setUser("admin")
     .setPassword("password")
+    .setConnectUsingSRVString(false)
     .setConnectionArguments("key=value;")
     .build();
 
@@ -45,6 +46,16 @@ public void testConfigConnectionString() {
                         VALID_CONFIG.getConnectionString());
   }
 
+  @Test
+  public void testConfigConnectionStringWithSRV() {
+    String connectionString = MongoDBConfigBuilder.builder(VALID_CONFIG)
+      .setConnectUsingSRVString(true)
+      .build()
+      .getConnectionString();
+
+    Assert.assertEquals("mongodb+srv://admin:password@localhost/admin.analytics?key=value;", connectionString);
+  }
+
   @Test
   public void testConfigConnectionStringNoCreds() {
     String connectionString = MongoDBConfigBuilder.builder(VALID_CONFIG)
@@ -56,6 +67,18 @@ public void testConfigConnectionStringNoCreds() {
     Assert.assertEquals("mongodb://localhost:27017/admin.analytics?key=value;", connectionString);
   }
 
+  @Test
+  public void testConfigConnectionStringWithSRVNoCreds() {
+    String connectionString = MongoDBConfigBuilder.builder(VALID_CONFIG)
+      .setConnectUsingSRVString(true)
+      .setUser(null)
+      .setPassword(null)
+      .build()
+      .getConnectionString();
+
+    Assert.assertEquals("mongodb+srv://localhost/admin.analytics?key=value;", connectionString);
+  }
+
   @Test
   public void testValidateValid() {
     VALID_CONFIG.validate();

src/test/java/io/cdap/plugin/batch/sink/MongoDBSinkConfigBuilder.java

@@ -49,6 +49,6 @@ public MongoDBSinkConfigBuilder setIdField(String idField) {
 
   public MongoDBBatchSink.MongoDBSinkConfig build() {
     return new MongoDBBatchSink.MongoDBSinkConfig(referenceName, host, port, database, collection, user, password,
-                                                  connectionArguments, idField);
+                                                  connectUsingSRVString, connectionArguments, idField);
   }
 }

src/test/java/io/cdap/plugin/batch/source/MongoDBSourceConfigBuilder.java

@@ -70,7 +70,7 @@ public MongoDBSourceConfigBuilder setAuthConnectionString(String authConnectionS
 
   public MongoDBBatchSource.MongoDBSourceConfig build() {
     return new MongoDBBatchSource.MongoDBSourceConfig(referenceName, host, port, database, collection, user, password,
-                                                      connectionArguments, schema, inputQuery, onError,
-                                                      authConnectionString);
+                                                      connectUsingSRVString, connectionArguments, schema, inputQuery,
+                                                      onError, authConnectionString);
   }
 }